package click.kamil.springstatemachineexporter.ast.common;

import click.kamil.springstatemachineexporter.analysis.model.LibraryHint;
import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import click.kamil.springstatemachineexporter.analysis.resolver.PropertyResolver;
import click.kamil.springstatemachineexporter.model.State;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.eclipse.jdt.core.dom.*;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class CodebaseContext {
    private final Map<String, CompilationUnit> classes = new HashMap<>();
    private final Map<String, Path> classPaths = new HashMap<>();
    private final Map<String, String> simpleNameToFqn = new HashMap<>();
    private final Set<String> ambiguousSimpleNames = new HashSet<>();
    private final ConstantResolver constantResolver = new ConstantResolver();
    private final PropertyResolver propertyResolver = new PropertyResolver();
    private Map<String, Map<String, String>> allProperties = new HashMap<>();
    private List<LibraryHint> libraryHints = new ArrayList<>();
    private final ObjectMapper objectMapper = new ObjectMapper();
    private Path projectRoot;

    public void setProjectRoot(Path root) {
        this.projectRoot = root;
    }

    public String getRelativePath(Path fullPath) {
        if (projectRoot != null && fullPath.isAbsolute() && fullPath.startsWith(projectRoot)) {
            return projectRoot.relativize(fullPath).toString();
        }
        return fullPath.getFileName().toString();
    }

    public String getRelativePath(String fqn) {
        Path path = classPaths.get(fqn);
        return path != null ? getRelativePath(path) : null;
    }

    public void loadLibraryHints(Path hintsFile) throws IOException {
        if (Files.exists(hintsFile)) {
            System.out.println("Loading hints from " + hintsFile.toAbsolutePath());
            this.libraryHints = objectMapper.readValue(hintsFile.toFile(), new TypeReference<List<LibraryHint>>() {});
            System.out.println("Loaded " + libraryHints.size() + " library hints");
        }
    }

    public List<LibraryHint> getLibraryHints() {
        return Collections.unmodifiableList(libraryHints);
    }

    private final Set<String> ignorePatterns = new HashSet<>(Set.of("**/test/**", "**/build/**", "**/node_modules/**", "**/out/**"));
    private String[] classpath = new String[0];
    private String[] sourcepath = new String[0];
    private boolean resolveBindings = false;

    public Map<String, Map<String, String>> getProperties() {
        return Collections.unmodifiableMap(allProperties);
    }

    public void setClasspath(List<String> classpath) {
        this.classpath = classpath.toArray(new String[0]);
    }

    public void setSourcepath(List<String> sourcepath) {
        this.sourcepath = sourcepath.toArray(new String[0]);
    }

    public void setResolveBindings(boolean resolveBindings) {
        this.resolveBindings = resolveBindings;
    }

    public void setActiveProfiles(List<String> profiles) {
        this.propertyResolver.setActiveProfiles(profiles);
    }

    public void scan(Path rootDir) throws IOException {
        scan(Set.of(rootDir), Collections.emptySet());
    }

    private final List<CompilationUnit> allCus = new ArrayList<>();

    private final Map<String, List<String>> interfaceToImpls = new HashMap<>();

    public void scan(Set<Path> rootDirs, Set<String> customIgnorePatterns) throws IOException {
        this.allProperties = propertyResolver.resolveAllProperties(rootDirs);

        Set<String> activeIgnore = new HashSet<>(ignorePatterns);
        activeIgnore.addAll(customIgnorePatterns);

        List<Path> javaFiles = FileUtils.findFilesWithExtension(rootDirs, ".java", activeIgnore);
        log.info("CodebaseContext found {} java files across {} root directories", javaFiles.size(), rootDirs.size());
        for (Path javaFile : javaFiles) {
            String source = Files.readString(javaFile);
            CompilationUnit cu = parse(source, javaFile.toAbsolutePath().toString());
            allCus.add(cu);
            
            String packageName = cu.getPackage() != null ? cu.getPackage().getName().getFullyQualifiedName() : "";
            for (Object type : cu.types()) {
                if (type instanceof TypeDeclaration td) {
                    indexType(td, packageName, javaFile);
                }
            }
        }
    }

    private void indexType(TypeDeclaration td, String parentFqn, Path javaFile) {
        String simpleName = td.getName().getIdentifier();
        String fqn = parentFqn.isEmpty() ? simpleName : parentFqn + "." + simpleName;
        
        classes.put(fqn, (CompilationUnit) td.getRoot());
        classPaths.put(fqn, javaFile);

        if (simpleNameToFqn.containsKey(simpleName)) {
            String existingFqn = simpleNameToFqn.get(simpleName);
            if (!existingFqn.equals(fqn)) {
                ambiguousSimpleNames.add(simpleName);
            }
        } else {
            simpleNameToFqn.put(simpleName, fqn);
        }

        // Inheritance Mapping
        for (Object itf : td.superInterfaceTypes()) {
            String itfName = itf.toString();
            interfaceToImpls.computeIfAbsent(itfName, k -> new ArrayList<>()).add(fqn);
        }

        // Recursively index nested types
        for (Object type : td.getTypes()) {
            if (type instanceof TypeDeclaration nestedTd) {
                indexType(nestedTd, fqn, javaFile);
            }
        }
    }

    public List<String> getImplementations(String interfaceName) {
        // Try direct match
        List<String> impls = interfaceToImpls.get(interfaceName);
        if (impls != null) return impls;
        
        // Try FQN match if input was simple name
        String fqn = simpleNameToFqn.get(interfaceName);
        if (fqn != null) return interfaceToImpls.getOrDefault(fqn, Collections.emptyList());
        
        return Collections.emptyList();
    }

    public State resolveState(Expression expr, CompilationUnit cu) {
        if (expr == null)
            return null;

        // 1. Check for constants
        String resolvedValue = constantResolver.resolve(expr, this);
        if (resolvedValue != null) {
            return State.of(expr.toString(), resolvedValue);
        }

        // 2. Check for @Value fields (SimpleName)
        if (expr instanceof SimpleName sn) {
            String fieldName = sn.getIdentifier();
            TypeDeclaration td = findEnclosingType(sn);
            if (td != null) {
                String value = findValueFromField(td, fieldName);
                if (value != null) {
                    return State.of(fieldName, value);
                }
            }
        }

        if (expr instanceof StringLiteral sl)
            return State.of(sl.getLiteralValue());

        String raw = expr.toString();
        if (expr instanceof QualifiedName qn) {
            String full = resolveQualifiedName(qn, cu);
            return State.of(raw, full);
        }
        if (expr instanceof SimpleName sn) {
            String full = resolveSimpleName(sn, cu);
            return State.of(raw, full);
        }
        return State.of(raw);
    }

    private String findValueFromField(TypeDeclaration td, String fieldName) {
        for (FieldDeclaration fd : td.getFields()) {
            for (Object fragObj : fd.fragments()) {
                if (fragObj instanceof VariableDeclarationFragment fragment) {
                    if (fragment.getName().getIdentifier().equals(fieldName)) {
                        // Check for @Value annotation
                        for (Object mod : fd.modifiers()) {
                            if (mod instanceof Annotation ann) {
                                String name = ann.getTypeName().getFullyQualifiedName();
                                if (name.endsWith("Value")) {
                                    return extractAnnotationValue(ann);
                                }
                            }
                        }
                    }
                }
            }
        }
        return null;
    }

    private String extractAnnotationValue(Annotation ann) {
        if (ann instanceof SingleMemberAnnotation sma) {
            return stripQuotes(sma.getValue().toString());
        } else if (ann instanceof NormalAnnotation na) {
            for (Object pairObj : na.values()) {
                MemberValuePair pair = (MemberValuePair) pairObj;
                if ("value".equals(pair.getName().getIdentifier())) {
                    return stripQuotes(pair.getValue().toString());
                }
            }
        }
        return null;
    }

    private String stripQuotes(String s) {
        if (s == null) return null;
        return s.replaceAll("^\"|\"$", "");
    }

    private TypeDeclaration findEnclosingType(ASTNode node) {
        ASTNode parent = node.getParent();
        while (parent != null && !(parent instanceof TypeDeclaration)) {
            parent = parent.getParent();
        }
        return (TypeDeclaration) parent;
    }

    private String resolveQualifiedName(QualifiedName qn, CompilationUnit cu) {
        String qualifier = qn.getQualifier().toString();
        String name = qn.getName().getIdentifier();
        TypeDeclaration td = getTypeDeclaration(qualifier, cu);
        if (td != null)
            return getFqn(td) + "." + name;
        return qn.getFullyQualifiedName();
    }

    private String resolveSimpleName(SimpleName sn, CompilationUnit cu) {
        String name = sn.getIdentifier();
        if (cu == null)
            return name;
        for (Object impObj : cu.imports()) {
            ImportDeclaration imp = (ImportDeclaration) impObj;
            String impName = imp.getName().getFullyQualifiedName();
            if (impName.endsWith("." + name))
                return impName;
        }
        return name;
    }

    private CompilationUnit parse(String source, String unitName) {
        ASTParser parser = ASTParser.newParser(AST.getJLSLatest());
        parser.setSource(source.toCharArray());
        parser.setKind(ASTParser.K_COMPILATION_UNIT);
        parser.setResolveBindings(resolveBindings);
        if (resolveBindings) {
            parser.setUnitName(unitName);
            parser.setEnvironment(classpath, sourcepath, null, true);
        }
        return (CompilationUnit) parser.createAST(null);
    }

    public TypeDeclaration getTypeDeclaration(String name) {
        // Try exact FQN match first
        CompilationUnit cu = classes.get(name);
        if (cu != null)
            return findTypeInCu(cu, name);

        // If it's a simple name, check if it's ambiguous
        if (ambiguousSimpleNames.contains(name)) {
            // System.err.println("Warning: Ambiguous simple name '" + name + "'. Provide FQN to resolve correctly.");
            return null; // Don't return a random class if it's ambiguous
        }

        // If not ambiguous, return the one we found during scan
        String fqn = simpleNameToFqn.get(name);
        if (fqn != null) {
            cu = classes.get(fqn);
            if (cu != null)
                return findTypeInCu(cu, fqn);
        }

        return null;
    }

    public TypeDeclaration getTypeDeclaration(String name, CompilationUnit contextCu) {
        if (name == null || name.isEmpty())
            return null;

        // 1. Check if it's already an FQN
        if (classes.containsKey(name))
            return getTypeDeclaration(name);

        // 2. Check imports in contextCu
        if (contextCu != null) {
            for (Object impObj : contextCu.imports()) {
                ImportDeclaration imp = (ImportDeclaration) impObj;
                String impName = imp.getName().getFullyQualifiedName();
                if (!imp.isStatic() && !imp.isOnDemand()) {
                    if (impName.endsWith("." + name))
                        return getTypeDeclaration(impName);
                }
            }

            // 3. Check same package
            String packageName = contextCu.getPackage() != null ? contextCu.getPackage().getName().getFullyQualifiedName() : "";
            String localFqn = packageName.isEmpty() ? name : packageName + "." + name;
            if (classes.containsKey(localFqn))
                return getTypeDeclaration(localFqn);

            // 4. Check on-demand imports (star imports)
            for (Object impObj : contextCu.imports()) {
                ImportDeclaration imp = (ImportDeclaration) impObj;
                if (imp.isOnDemand()) {
                    String starFqn = imp.getName().getFullyQualifiedName() + "." + name;
                    if (classes.containsKey(starFqn))
                        return getTypeDeclaration(starFqn);
                }
            }

            // 5. Fallback to java.lang (common)
            String langFqn = "java.lang." + name;
            if (classes.containsKey(langFqn))
                return getTypeDeclaration(langFqn);
        }

        // 6. Last resort: global simple name match (might be ambiguous but better than nothing)
        return getTypeDeclaration(name);
    }

    public String getFqn(TypeDeclaration td) {
        StringBuilder sb = new StringBuilder(td.getName().getIdentifier());
        ASTNode current = td.getParent();
        while (current instanceof TypeDeclaration parent) {
            sb.insert(0, parent.getName().getIdentifier() + ".");
            current = parent.getParent();
        }

        CompilationUnit cu = (CompilationUnit) td.getRoot();
        String packageName = cu.getPackage() != null ? cu.getPackage().getName().getFullyQualifiedName() : "";
        String fqn = sb.toString();
        return packageName.isEmpty() ? fqn : packageName + "." + fqn;
    }

    private TypeDeclaration findTypeInCu(CompilationUnit cu, String fqn) {
        for (Object type : cu.types()) {
            if (type instanceof TypeDeclaration td) {
                TypeDeclaration found = findNestedType(td, fqn);
                if (found != null) return found;
            }
        }
        return null;
    }

    private TypeDeclaration findNestedType(TypeDeclaration td, String targetFqn) {
        if (getFqn(td).equals(targetFqn)) return td;
        for (TypeDeclaration nested : td.getTypes()) {
            TypeDeclaration found = findNestedType(nested, targetFqn);
            if (found != null) return found;
        }
        return null;
    }

    public Path getPath(String className) {
        return classPaths.get(className);
    }

    public List<TypeDeclaration> findEntryPointClasses(List<String> targetAnnotations) {
        List<TypeDeclaration> entryPoints = new ArrayList<>();
        for (CompilationUnit cu : classes.values()) {
            for (Object type : cu.types()) {
                if (type instanceof TypeDeclaration td) {
                    if (isEntryPoint(td, targetAnnotations)) {
                        entryPoints.add(td);
                    }
                }
            }
        }
        return entryPoints;
    }

    private boolean isEntryPoint(TypeDeclaration td, List<String> targetAnnotations) {
        if (td.isInterface())
            return false;

        boolean isAbstract = false;
        // Has annotation
        for (Object modifier : td.modifiers()) {
            if (modifier instanceof Annotation annotation) {
                String annotationName = annotation.getTypeName().getFullyQualifiedName();
                String simpleAnnotationName = annotationName.contains(".")
                        ? annotationName.substring(annotationName.lastIndexOf('.') + 1)
                        : annotationName;

                if (targetAnnotations.contains(annotationName) || targetAnnotations.contains(simpleAnnotationName)) {
                    return true;
                }
            }
            if (modifier instanceof Modifier m && m.isAbstract()) {
                isAbstract = true;
            }
        }

        // If it's a leaf class (not abstract) and extends StateMachineConfigurerAdapter indirectly
        if (!isAbstract && extendsStateMachineConfigurerAdapter(td)) {
            return true;
        }

        return false;
    }

    public boolean extendsStateMachineConfigurerAdapter(TypeDeclaration td) {
        return extendsStateMachineConfigurerAdapter(td, new HashSet<>());
    }

    private boolean extendsStateMachineConfigurerAdapter(TypeDeclaration td, Set<String> visited) {
        String fqn = getFqn(td);
        if (visited.contains(fqn)) {
            return false;
        }
        visited.add(fqn);

        Set<String> knownAdapters = Set.of("EnumStateMachineConfigurerAdapter", "StateMachineConfigurerAdapter", "StateMachineConfigurer");
        CompilationUnit cu = (CompilationUnit) td.getRoot();

        // Check superclass
        Type superclass = td.getSuperclassType();
        if (superclass != null) {
            String superclassName = getSimpleName(superclass);
            if (knownAdapters.contains(superclassName))
                return true;
            TypeDeclaration superTd = getTypeDeclaration(superclassName, cu);
            if (superTd != null && extendsStateMachineConfigurerAdapter(superTd, visited))
                return true;
        }

        // Check interfaces
        for (Object interfaceTypeObj : td.superInterfaceTypes()) {
            if (interfaceTypeObj instanceof Type interfaceType) {
                String interfaceName = getSimpleName(interfaceType);
                if (knownAdapters.contains(interfaceName))
                    return true;
                TypeDeclaration interfaceTd = getTypeDeclaration(interfaceName, cu);
                if (interfaceTd != null && extendsStateMachineConfigurerAdapter(interfaceTd, visited))
                    return true;
            }
        }

        return false;
    }

    public List<MethodDeclaration> findBeanMethodsReturning(Set<String> targetReturnTypes) {
        List<MethodDeclaration> beanMethods = new ArrayList<>();
        for (CompilationUnit cu : classes.values()) {
            for (Object type : cu.types()) {
                if (type instanceof TypeDeclaration td) {
                    for (MethodDeclaration method : td.getMethods()) {
                        if (isBeanMethodReturning(method, targetReturnTypes)) {
                            beanMethods.add(method);
                        }
                    }
                }
            }
        }
        return beanMethods;
    }

    private boolean isBeanMethodReturning(MethodDeclaration method, Set<String> targetReturnTypes) {
        // Check return type
        Type returnType = method.getReturnType2();
        if (returnType == null)
            return false;
        String typeName = AstUtils.extractSimpleTypeName(returnType);
        if (!targetReturnTypes.contains(typeName))
            return false;

        // Check for @Bean annotation
        for (Object modifier : method.modifiers()) {
            if (modifier instanceof Annotation annotation) {
                String annotationName = annotation.getTypeName().getFullyQualifiedName();
                if ("Bean".equals(annotationName) || "org.springframework.context.annotation.Bean".equals(annotationName)) {
                    return true;
                }
            }
        }
        return false;
    }

    private String getSimpleName(Type type) {
        return AstUtils.extractSimpleTypeName(type);
    }

    public Collection<CompilationUnit> getCompilationUnits() {
        return Collections.unmodifiableCollection(allCus);
    }

    public Collection<TypeDeclaration> getTypeDeclarations() {
        List<TypeDeclaration> types = new ArrayList<>();
        for (CompilationUnit cu : allCus) {
            for (Object typeObj : cu.types()) {
                if (typeObj instanceof TypeDeclaration td) {
                    types.add(td);
                    // Handle nested types
                    types.addAll(getNestedTypes(td));
                }
            }
        }
        return types;
    }

    private List<TypeDeclaration> getNestedTypes(TypeDeclaration td) {
        List<TypeDeclaration> nested = new ArrayList<>();
        for (TypeDeclaration inner : td.getTypes()) {
            nested.add(inner);
            nested.addAll(getNestedTypes(inner));
        }
        return nested;
    }
}
