package click.kamil.file_combine;

import click.kamil.file_combine.common.FileUtils;
import org.eclipse.jdt.core.JavaCore;
import org.eclipse.jdt.core.compiler.IProblem;
import org.eclipse.jdt.core.dom.*;
import org.eclipse.jdt.core.dom.rewrite.ASTRewrite;
import org.eclipse.jface.text.Document;
import org.eclipse.text.edits.TextEdit;
import picocli.CommandLine;
import picocli.CommandLine.Command;
import picocli.CommandLine.Option;
import picocli.CommandLine.Parameters;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.concurrent.Callable;

@Command(name = "java-combine", mixinStandardHelpOptions = true, version = "1.0",
        description = "Combines multiple Java source files into a single file.")
public class JavaConcatenator implements Callable<Integer> {

    @Parameters(index = "0..*", description = "Directories to search for Java files.")
    private List<Path> sourceDirs;

    @Option(names = {"-o", "--output"}, description = "Output file path.", defaultValue = "./Combined.java")
    private Path outputFile;

    @Option(names = {"-p", "--package"}, description = "Target package name.", defaultValue = "click.kamil")
    private String targetPackage;

    public static void main(String[] args) {
        int exitCode = new CommandLine(new JavaConcatenator()).execute(args);
        System.exit(exitCode);
    }

    @Override
    public Integer call() throws IOException {
        if (sourceDirs == null || sourceDirs.isEmpty()) {
            System.err.println("At least one source directory must be specified.");
            return 1;
        }

        List<Path> javaFiles = FileUtils.findFilesWithExtension(new HashSet<>(sourceDirs), ".java");

        if (javaFiles.isEmpty()) {
            System.out.println("No Java files found in specified directories.");
            return 0;
        }

        Set<String> internalPackages = new HashSet<>();
        Set<String> internalClasses = new HashSet<>();
        List<FileParseResult> results = new ArrayList<>();

        for (Path javaFile : javaFiles) {
            String source = Files.readString(javaFile);
            FileParseResult result = parseAndTransform(source);
            if (result.packageName != null) {
                internalPackages.add(result.packageName);
            }
            internalClasses.addAll(result.topLevelTypes);
            results.add(result);
        }

        Set<String> uniqueImports = new TreeSet<>();
        for (FileParseResult result : results) {
            for (String imp : result.imports) {
                String normalizedImport = imp.replaceFirst("^import\\s+static\\s+", "")
                        .replaceFirst("^import\\s+", "")
                        .replaceFirst(";$", "");

                boolean isInternal = false;
                for (String pkg : internalPackages) {
                    if (normalizedImport.startsWith(pkg + ".")) {
                        isInternal = true;
                        break;
                    }
                }

                if (!isInternal && (targetPackage == null || !normalizedImport.startsWith(targetPackage + "."))) {
                    uniqueImports.add(imp);
                }
            }
        }

        StringBuilder combinedSource = new StringBuilder();

        // Add package declaration
        if (targetPackage != null && !targetPackage.isEmpty()) {
            combinedSource.append("package ").append(targetPackage).append(";\n\n");
        }

        // Append imports
        for (String imp : uniqueImports) {
            combinedSource.append(imp).append("\n");
        }
        combinedSource.append("\n");

        // Append all class bodies
        for (FileParseResult result : results) {
            combinedSource.append(result.classBody).append("\n");
        }

        if (outputFile.getParent() != null) {
            Files.createDirectories(outputFile.getParent());
        }
        Files.writeString(outputFile, combinedSource.toString());

        System.out.println("Combined Java source written to: " + outputFile);
        return 0;
    }

    private static FileParseResult parseAndTransform(String source) {
        ASTParser parser = ASTParser.newParser(AST.getJLSLatest());
        parser.setSource(source.toCharArray());
        parser.setKind(ASTParser.K_COMPILATION_UNIT);
        parser.setResolveBindings(false);

        Map<String, String> options = JavaCore.getOptions();
        JavaCore.setComplianceOptions(JavaCore.VERSION_21, options);
        parser.setCompilerOptions(options);

        CompilationUnit cu = (CompilationUnit) parser.createAST(null);
        
        if (cu.getProblems().length > 0) {
            for (IProblem problem : cu.getProblems()) {
                if (problem.isError()) {
                    System.err.println("Error parsing AST: " + problem.getMessage() + " at line " + problem.getSourceLineNumber());
                }
            }
        }

        AST ast = cu.getAST();
        ASTRewrite rewrite = ASTRewrite.create(ast);

        String packageName = cu.getPackage() != null ? cu.getPackage().getName().getFullyQualifiedName() : null;

        // Collect imports as strings
        List<String> imports = new ArrayList<>();
        for (Object obj : cu.imports()) {
            ImportDeclaration imp = (ImportDeclaration) obj;
            imports.add(imp.toString().trim());
            rewrite.remove(imp, null); // Remove from AST
        }

        // Remove package declaration
        if (cu.getPackage() != null) {
            rewrite.remove(cu.getPackage(), null);
        }

        List<String> topLevelTypes = new ArrayList<>();

        // Remove 'public' from top-level type declarations
        cu.accept(new ASTVisitor() {
            @Override
            public boolean visit(TypeDeclaration node) {
                if (node.getParent() instanceof CompilationUnit) {
                    topLevelTypes.add(node.getName().getIdentifier());
                    removePublicModifier(node.modifiers(), rewrite);
                }
                return true;
            }

            @Override
            public boolean visit(EnumDeclaration node) {
                if (node.getParent() instanceof CompilationUnit) {
                    topLevelTypes.add(node.getName().getIdentifier());
                    removePublicModifier(node.modifiers(), rewrite);
                }
                return true;
            }

            @Override
            public boolean visit(AnnotationTypeDeclaration node) {
                if (node.getParent() instanceof CompilationUnit) {
                    topLevelTypes.add(node.getName().getIdentifier());
                    removePublicModifier(node.modifiers(), rewrite);
                }
                return true;
            }

            @Override
            public boolean visit(RecordDeclaration node) {
                if (node.getParent() instanceof CompilationUnit) {
                    topLevelTypes.add(node.getName().getIdentifier());
                    removePublicModifier(node.modifiers(), rewrite);
                }
                return true;
            }

            private void removePublicModifier(List modifiers, ASTRewrite rewrite) {
                for (Object mod : modifiers) {
                    if (mod instanceof Modifier) {
                        Modifier m = (Modifier) mod;
                        if (m.isPublic()) {
                            rewrite.remove(m, null);
                        }
                    }
                }
            }
        });

        try {
            Document document = new Document(source);
            TextEdit edits = rewrite.rewriteAST(document, null);
            edits.apply(document);
            String transformedSource = document.get();
            return new FileParseResult(packageName, topLevelTypes, imports, transformedSource);
        } catch (Exception e) {
            throw new RuntimeException("Failed to apply AST edits", e);
        }
    }

    private static class FileParseResult {
        String packageName;
        List<String> topLevelTypes;
        List<String> imports;
        String classBody;

        public FileParseResult(String packageName, List<String> topLevelTypes, List<String> imports, String classBody) {
            this.packageName = packageName;
            this.topLevelTypes = topLevelTypes;
            this.imports = imports;
            this.classBody = classBody;
        }
    }
}
