// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

import registerDebug from "debug";
import { Grammar } from "./grammarTypes.js";
import { NFA } from "./nfa.js";
import {
    matchNFAWithIndex,
    buildFirstTokenIndex,
    NFAMatchResult,
    type FirstTokenIndex,
} from "./nfaInterpreter.js";
import { compileGrammarToNFA } from "./nfaCompiler.js";
import { loadGrammarRulesNoThrow } from "./grammarLoader.js";
import { mergeGrammarRules } from "./grammarMerger.js";
import { globalEntityRegistry } from "./entityRegistry.js";

const debug = registerDebug("typeagent:actionGrammar:registry");

/**
 * Agent Grammar Registry
 *
 * Manages in-memory grammar representations for each active agent.
 * Supports dynamic addition of generated rules for cache hit optimization.
 */

/**
 * Represents a single agent's grammar with its compiled NFA
 */
export class AgentGrammar {
    private grammar: Grammar;
    private nfa: NFA;
    private firstTokenIndex: FirstTokenIndex;
    private ruleCount: number;
    private readonly baseGrammar: Grammar; // Store original grammar for reset
    private readonly baseNFA: NFA; // Store original NFA for reset
    private readonly baseFirstTokenIndex: FirstTokenIndex;

    constructor(
        public readonly agentId: string,
        grammar: Grammar,
        nfa: NFA,
    ) {
        this.grammar = grammar;
        this.nfa = nfa;
        this.firstTokenIndex = buildFirstTokenIndex(nfa);
        this.ruleCount = grammar.rules.length;
        // Store deep copy of base grammar for reset capability
        this.baseGrammar = JSON.parse(JSON.stringify(grammar));
        this.baseNFA = nfa; // NFA is immutable, safe to share reference
        this.baseFirstTokenIndex = this.firstTokenIndex;
    }

    /**
     * Add dynamically generated rules to this agent's grammar
     *
     * Rules are typically generated by grammarGenerator.ts from user requests.
     * They are merged as alternatives to enable cache hits on similar requests.
     *
     * @param agrText Grammar rules in .agr format (from grammarGenerator)
     * @param checkedVariables Optional set of variable names with validation (checked_wildcard paramSpec)
     * @returns Success status and any errors
     */
    addGeneratedRules(
        agrText: string,
        checkedVariables?: Set<string>,
    ): {
        success: boolean;
        errors: string[];
        unresolvedEntities?: string[];
    } {
        const errors: string[] = [];

        // Prepend entity declarations as source-less imports so generated rules
        // can reference entity types like CalendarDate, CalendarTime, Ordinal, etc.
        // Collect from both the existing grammar and the global registry.
        const entityNames = new Set<string>(this.grammar.entities || []);
        for (const name of globalEntityRegistry.getEntityNames()) {
            entityNames.add(name);
        }
        let fullAgrText = agrText;
        if (entityNames.size > 0) {
            fullAgrText = `import { ${[...entityNames].join(", ")} };\n${agrText}`;
        }

        // Parse the generated rules
        const newGrammar = loadGrammarRulesNoThrow(
            `<generated-${this.agentId}>`,
            fullAgrText,
            errors,
        );

        if (!newGrammar) {
            return {
                success: false,
                errors: [
                    `Failed to parse generated rules: ${errors.join(", ")}`,
                ],
            };
        }

        // Add checked variables if provided
        if (checkedVariables && checkedVariables.size > 0) {
            newGrammar.checkedVariables = checkedVariables;
        }

        // Validate entity references
        const unresolvedEntities = this.validateEntityReferences(newGrammar);
        if (unresolvedEntities.length > 0) {
            return {
                success: false,
                errors: [
                    "Generated rules reference unresolved entities. Register entities before adding rules.",
                ],
                unresolvedEntities,
            };
        }

        // Merge into existing grammar
        const mergedGrammar = mergeGrammarRules(this.grammar, newGrammar.rules);

        // Merge entity declarations
        if (newGrammar.entities) {
            const existingEntities = new Set(this.grammar.entities || []);
            for (const entity of newGrammar.entities) {
                existingEntities.add(entity);
            }
            mergedGrammar.entities = Array.from(existingEntities);
        }

        // Merge checked variables
        if (newGrammar.checkedVariables || this.grammar.checkedVariables) {
            const existingChecked = new Set(
                this.grammar.checkedVariables || [],
            );
            if (newGrammar.checkedVariables) {
                for (const varName of newGrammar.checkedVariables) {
                    existingChecked.add(varName);
                }
            }
            mergedGrammar.checkedVariables = existingChecked;
        }

        // Recompile NFA
        try {
            const newNFA = compileGrammarToNFA(
                mergedGrammar,
                `${this.agentId}-cache`,
            );
            const previousRuleCount = this.ruleCount;
            this.grammar = mergedGrammar;
            this.nfa = newNFA;
            this.firstTokenIndex = buildFirstTokenIndex(newNFA);
            this.ruleCount = mergedGrammar.rules.length;

            // Log the newly added grammar rules for debugging/testing
            const newRulesCount = this.ruleCount - previousRuleCount;
            if (newRulesCount > 0) {
                debug(
                    `Added ${newRulesCount} new rule(s) to agent '%s': %s (total: %d)`,
                    this.agentId,
                    agrText,
                    this.ruleCount,
                );
            }

            return { success: true, errors: [] };
        } catch (error) {
            return {
                success: false,
                errors: [
                    `Failed to compile merged grammar: ${error instanceof Error ? error.message : String(error)}`,
                ],
            };
        }
    }

    /**
     * Reset this agent's grammar to its original base state, removing all dynamically added rules
     *
     * This is used when creating a new construction store to start fresh.
     */
    resetToBase(): void {
        debug(
            `Resetting agent '%s' to base grammar (removing %d dynamic rules)`,
            this.agentId,
            this.ruleCount - this.baseGrammar.rules.length,
        );
        this.grammar = JSON.parse(JSON.stringify(this.baseGrammar));
        this.nfa = this.baseNFA;
        this.firstTokenIndex = this.baseFirstTokenIndex;
        this.ruleCount = this.baseGrammar.rules.length;
    }

    /**
     * Match tokens against this agent's grammar
     */
    match(tokens: string[]): NFAMatchResult {
        return matchNFAWithIndex(this.nfa, this.firstTokenIndex, tokens);
    }

    /**
     * Get current grammar
     */
    getGrammar(): Grammar {
        return this.grammar;
    }

    /**
     * Get current NFA
     */
    getNFA(): NFA {
        return this.nfa;
    }

    /**
     * Get statistics about this agent's grammar
     */
    getStats() {
        return {
            agentId: this.agentId,
            ruleCount: this.ruleCount,
            stateCount: this.nfa.states.length,
            transitionCount: this.nfa.states.reduce(
                (sum, s) => sum + s.transitions.length,
                0,
            ),
        };
    }

    private validateEntityReferences(grammar: Grammar): string[] {
        const unresolved = new Set<string>();
        const builtInTypes = new Set(["string", "number", "wildcard"]);

        // Extract known type names from the existing NFA's wildcard transitions
        // These are valid type references since they're already used in the base grammar
        const knownNFATypes = this.extractTypeNamesFromNFA();

        // Check declared entities
        if (grammar.entities) {
            for (const entity of grammar.entities) {
                if (!globalEntityRegistry.hasEntity(entity)) {
                    unresolved.add(entity);
                }
            }
        }

        // Check wildcard types in rules
        for (const rule of grammar.rules) {
            this.checkPartsForEntities(
                rule.parts,
                unresolved,
                builtInTypes,
                knownNFATypes,
            );
        }

        return Array.from(unresolved);
    }

    /**
     * Extract all type names used in wildcard transitions from the NFA
     * These represent valid type references from the base grammar
     */
    private extractTypeNamesFromNFA(): Set<string> {
        const typeNames = new Set<string>();
        for (const state of this.nfa.states) {
            for (const transition of state.transitions) {
                if (transition.type === "wildcard" && transition.typeName) {
                    typeNames.add(transition.typeName);
                }
            }
        }
        return typeNames;
    }

    private checkPartsForEntities(
        parts: any[],
        unresolved: Set<string>,
        builtInTypes: Set<string>,
        knownNFATypes: Set<string>,
    ): void {
        for (const part of parts) {
            if (part.type === "wildcard" && part.typeName) {
                // Type is valid if it's:
                // 1. A built-in type (string, number, wildcard)
                // 2. Registered in the global entity registry
                // 3. A type name already used in the NFA (from base grammar rules)
                if (
                    !builtInTypes.has(part.typeName) &&
                    !globalEntityRegistry.hasEntity(part.typeName) &&
                    !knownNFATypes.has(part.typeName)
                ) {
                    unresolved.add(part.typeName);
                }
            } else if (part.type === "rules") {
                for (const nestedRule of part.rules) {
                    this.checkPartsForEntities(
                        nestedRule.parts,
                        unresolved,
                        builtInTypes,
                        knownNFATypes,
                    );
                }
            }
        }
    }
}

/**
 * Result from matching across all agents
 */
export interface AgentMatchResult {
    matched: boolean;
    agentId?: string | undefined; // Which agent matched
    actionValue?: any | undefined; // Evaluated action object from the matched rule
    // Priority counts for sorting matches
    fixedStringPartCount?: number | undefined;
    checkedWildcardCount?: number | undefined;
    uncheckedWildcardCount?: number | undefined;
    // Debugging info
    attemptedAgents?: string[] | undefined;
    tokensConsumed?: number | undefined;
    debugSlotMap?: Map<string, number> | undefined; // Variable name -> slot index (debug only)
}

/**
 * Registry for managing agent grammars
 *
 * Maintains in-memory grammar representations for each active agent.
 * Supports dynamic rule addition and multi-agent matching.
 */
export class AgentGrammarRegistry {
    private agents = new Map<string, AgentGrammar>();

    /**
     * Register an agent with its grammar
     *
     * @param agentId Unique identifier for the agent
     * @param grammar The agent's grammar
     * @param nfa Optional pre-compiled NFA (will compile if not provided)
     * @returns The registered AgentGrammar
     */
    registerAgent(agentId: string, grammar: Grammar, nfa?: NFA): AgentGrammar {
        const compiledNFA = nfa ?? compileGrammarToNFA(grammar, agentId);
        const agentGrammar = new AgentGrammar(agentId, grammar, compiledNFA);
        this.agents.set(agentId, agentGrammar);
        return agentGrammar;
    }

    /**
     * Register an agent from .agr grammar text
     *
     * @param agentId Unique identifier for the agent
     * @param agrText Grammar rules in .agr format
     * @returns Success status with the AgentGrammar or errors
     */
    registerAgentFromText(
        agentId: string,
        agrText: string,
    ): {
        success: boolean;
        agentGrammar?: AgentGrammar;
        errors?: string[];
    } {
        const errors: string[] = [];

        const grammar = loadGrammarRulesNoThrow(
            `${agentId}.agr`,
            agrText,
            errors,
        );

        if (!grammar) {
            return {
                success: false,
                errors: [`Failed to load grammar: ${errors.join(", ")}`],
            };
        }

        try {
            const agentGrammar = this.registerAgent(agentId, grammar);
            return { success: true, agentGrammar };
        } catch (error) {
            return {
                success: false,
                errors: [
                    `Failed to compile grammar: ${error instanceof Error ? error.message : String(error)}`,
                ],
            };
        }
    }

    /**
     * Unregister an agent
     */
    unregisterAgent(agentId: string): boolean {
        return this.agents.delete(agentId);
    }

    /**
     * Get an agent's grammar
     */
    getAgent(agentId: string): AgentGrammar | undefined {
        return this.agents.get(agentId);
    }

    /**
     * Get all registered agent IDs
     */
    getAllAgentIds(): string[] {
        return Array.from(this.agents.keys());
    }

    /**
     * Add dynamically generated rules to an agent's grammar
     *
     * This is the primary method for cache optimization. When grammarGenerator
     * creates rules from user requests, they are added here to enable cache hits.
     *
     * @param agentId The agent to add rules to
     * @param agrText Generated rules in .agr format
     * @returns Success status and any errors
     */
    addGeneratedRules(
        agentId: string,
        agrText: string,
    ): {
        success: boolean;
        errors: string[];
        unresolvedEntities?: string[];
    } {
        let agent = this.agents.get(agentId);
        if (!agent) {
            // Auto-register with empty grammar so dynamic rules accumulate even
            // for agents that have no pre-compiled base grammar file.
            debug(
                `Agent '%s' not in registry; creating empty grammar for dynamic rule addition`,
                agentId,
            );
            try {
                agent = this.registerAgent(agentId, { rules: [] });
            } catch (e) {
                return {
                    success: false,
                    errors: [
                        `Agent '${agentId}' not found and could not create empty grammar: ${e instanceof Error ? e.message : String(e)}`,
                    ],
                };
            }
        }

        return agent.addGeneratedRules(agrText);
    }

    /**
     * Match tokens against all registered agents
     *
     * Tries each agent's grammar in registration order until a match is found.
     * This enables the cache to work across multiple agents.
     *
     * @param tokens Input tokens to match
     * @param agentIds Optional list of specific agents to try (in order)
     * @returns Match result with the agent that matched
     */
    matchAcrossAgents(tokens: string[], agentIds?: string[]): AgentMatchResult {
        const attemptedAgents: string[] = [];
        const agentsToTry = agentIds
            ? agentIds
                  .map((id) => ({ id, agent: this.agents.get(id) }))
                  .filter((a) => a.agent !== undefined)
            : Array.from(this.agents.entries()).map(([id, agent]) => ({
                  id,
                  agent,
              }));

        for (const { id, agent } of agentsToTry) {
            attemptedAgents.push(id);
            const result = agent!.match(tokens);

            if (result.matched) {
                return {
                    matched: true,
                    agentId: id,
                    actionValue: result.actionValue,
                    fixedStringPartCount: result.fixedStringPartCount,
                    checkedWildcardCount: result.checkedWildcardCount,
                    uncheckedWildcardCount: result.uncheckedWildcardCount,
                    attemptedAgents,
                    tokensConsumed: result.tokensConsumed,
                    debugSlotMap: result.debugSlotMap,
                };
            }
        }

        // No match found
        return {
            matched: false,
            attemptedAgents,
            tokensConsumed: 0,
        };
    }

    /**
     * Get all registered agent IDs
     */
    getAgentIds(): string[] {
        return Array.from(this.agents.keys());
    }

    /**
     * Get statistics about all agents
     */
    getAllStats() {
        return Array.from(this.agents.values()).map((agent) =>
            agent.getStats(),
        );
    }

    /**
     * Clear all registered agents
     */
    clear(): void {
        this.agents.clear();
    }

    /**
     * Reset all agents to their base grammars, removing dynamically added rules
     *
     * This is useful when clearing the cache to start fresh while keeping
     * the original static grammar rules intact.
     */
    resetAllToBase(): void {
        debug(`Resetting %d agent(s) to base grammar`, this.agents.size);
        for (const agent of this.agents.values()) {
            agent.resetToBase();
        }
    }

    /**
     * Placeholder for future DFA compilation
     *
     * In the future, this method will asynchronously compile all agent NFAs
     * into a single optimized DFA for faster matching.
     *
     * @returns Promise that resolves when compilation is complete
     */
    async compileToDFA(): Promise<{
        success: boolean;
        stats?: { stateCount: number; transitionCount: number };
    }> {
        // TODO: Implement async DFA compilation
        // This would:
        // 1. Merge all agent NFAs
        // 2. Convert the combined NFA to DFA
        // 3. Minimize the DFA
        // 4. Replace individual NFAs with DFA for matching
        return {
            success: false,
        };
    }
}

/**
 * Global agent grammar registry instance
 *
 * Use this for managing agent grammars across the application.
 */
export const globalAgentGrammarRegistry = new AgentGrammarRegistry();
