import { type Module, inject } from 'langium';
import { createDefaultModule, createDefaultSharedModule, type DefaultSharedModuleContext, type LangiumServices, type LangiumSharedServices, type PartialLangiumServices } from 'langium/lsp';
import { IsaGeneratedModule, IsaGeneratedSharedModule } from './generated/module.js';
import { IsaValidator, registerValidationChecks } from './isa-validator.js';
import { IsaCompletionProvider } from './isa-completion-provider.js';
import { IsaScopeProvider } from './isa-scope-provider.js';
import { IsaDocumentProcessor } from './isa-document-processor.js';

/**
 * Declaration of custom services - add your own service classes here.
 */
export type IsaAddedServices = {
    validation: {
        IsaValidator: IsaValidator
    },
    completion: {
        CompletionProvider: IsaCompletionProvider
    },
    references: {
        ScopeProvider: IsaScopeProvider
    },
    workspace: {
        DocumentProcessor: IsaDocumentProcessor
    }
}

/**
 * Union of Langium default services and your custom services - use this as constructor parameter
 * of custom service classes.
 */
export type IsaServices = LangiumServices & IsaAddedServices

/**
 * Dependency injection module that overrides Langium default services and contributes the
 * declared custom services. The Langium defaults can be partially specified to override only
 * selected services, while the custom services must be fully specified.
 */
export const IsaModule: Module<IsaServices, PartialLangiumServices & IsaAddedServices> = {
    validation: {
        IsaValidator: () => new IsaValidator()
    },
    completion: {
        CompletionProvider: (services) => new IsaCompletionProvider(services)
    },
    references: {
        ScopeProvider: (services) => new IsaScopeProvider(services)
    },
    workspace: {
        DocumentProcessor: (services) => new IsaDocumentProcessor(services)
    }
};

/**
 * Create the full set of services required by Langium.
 *
 * First inject the shared services by merging two modules:
 *  - Langium default shared services
 *  - Services generated by langium-cli
 *
 * Then inject the language-specific services by merging three modules:
 *  - Langium default language-specific services
 *  - Services generated by langium-cli
 *  - Services specified in this file
 *
 * @param context Optional module context with the LSP connection
 * @returns An object wrapping the shared services and the language-specific services
 */
export function createIsaServices(context: DefaultSharedModuleContext): {
    shared: LangiumSharedServices,
    Isa: IsaServices
} {
    const shared = inject(
        createDefaultSharedModule(context),
        IsaGeneratedSharedModule
    );
    const Isa = inject(
        createDefaultModule({ shared }),
        IsaGeneratedModule,
        IsaModule
    );
    shared.ServiceRegistry.register(Isa);
    registerValidationChecks(Isa);
    
    // Hook into document builder to process includes BEFORE documents are fully built
    // This ensures included files are available when resolving cross-references during linking
    const documentProcessor = Isa.workspace.DocumentProcessor;
    let isProcessingIncludes = false; // Flag to prevent recursive calls
    const originalBuild = shared.workspace.DocumentBuilder.build.bind(shared.workspace.DocumentBuilder);
    shared.workspace.DocumentBuilder.build = async (documents, options, cancelToken) => {
        // If we're already processing includes (called from document processor), just build normally
        // But still wrap in error handling to prevent crashes
        if (isProcessingIncludes) {
            try {
                await Promise.race([
                    originalBuild(documents, options, cancelToken),
                    new Promise((_, reject) => setTimeout(() => reject(new Error('Parsing timeout')), 5000))
                ]);
                return;
            } catch (error: any) {
                // Handle recursion errors gracefully even in include processing
                if (error instanceof RangeError || 
                    (error.message && (error.message.includes('Maximum call stack') || error.message.includes('timeout')))) {
                    console.warn('Chevrotain recursion detected during include processing. Returning partial results.');
                    return; // Return void to allow partial results
                }
                throw error;
            }
        }
        
        // Step 1: Build documents with parsing only (no linking/validation) to extract includes
        // We need to parse first to get the AST and extract include directives
        // Note: We can't skip linking completely, but we'll rebuild after processing includes
        // Wrap in try-catch with timeout to handle Chevrotain recursion errors with hex values in RTL expressions
        try {
            await Promise.race([
                originalBuild(documents, { ...options, validation: false }, cancelToken),
                new Promise((_, reject) => setTimeout(() => reject(new Error('Parsing timeout')), 5000))
            ]);
        } catch (error: any) {
            // If we hit recursion (Maximum call stack or RangeError), it's likely due to hex values in RTL expressions
            // This is a known Chevrotain limitation - the file is still valid (works in Python)
            // Log the error but don't fail completely - allow partial parsing
            if (error instanceof RangeError || 
                (error.message && (error.message.includes('Maximum call stack') || error.message.includes('timeout')))) {
                console.warn('Chevrotain recursion detected during parsing (likely hex values in RTL expressions). File may be partially parsed.');
                // Try to continue with what we have - documents may have partial parse results
                // The error will be in the document diagnostics, which is acceptable
            } else {
                throw error;
            }
        }
        
        // Step 2: Process includes for all documents (this loads included files)
        // We need to do this AFTER parsing but BEFORE full linking
        isProcessingIncludes = true;
        try {
            for (const doc of documents) {
                try {
                    await documentProcessor.processDocument(doc);
                } catch (error) {
                    // Let validation handle errors
                }
            }
        } finally {
            isProcessingIncludes = false;
        }
        
        // Step 3: Now rebuild all documents (original + included) with full linking
        // At this point, included files are loaded, so cross-references can be resolved
        // First, invalidate all documents to force re-linking
        const allDocuments = Array.from(shared.workspace.LangiumDocuments.all);
        
        // Invalidate all documents to force re-linking with the new included files
        for (const doc of allDocuments) {
            shared.workspace.LangiumDocuments.invalidateDocument(doc.uri);
        }
        
        // Now rebuild all documents with full linking and validation
        // Wrap in try-catch with timeout to handle Chevrotain recursion errors
        try {
            await Promise.race([
                originalBuild(allDocuments, options, cancelToken),
                new Promise((_, reject) => setTimeout(() => reject(new Error('Parsing timeout')), 5000))
            ]);
            return;
        } catch (error: any) {
            // If we hit recursion during final build (RangeError or Maximum call stack), log and return what we have
            if (error instanceof RangeError || 
                (error.message && (error.message.includes('Maximum call stack') || error.message.includes('timeout')))) {
                console.warn('Chevrotain recursion detected during final build (likely hex values in RTL expressions). Returning partial results.');
                // Return void - documents will have diagnostics indicating the issue
                // This is better than crashing the language server
                return;
            }
            throw error;
        }
    };
    
    if (!context.connection) {
        // We don't run inside a language server
        // Therefore, initialize the configuration provider instantly
        shared.workspace.ConfigurationProvider.initialized({});
    }
    return { shared, Isa };
}
