import { App, MarkdownView, TFile, TFolder, WorkspaceLeaf, normalizePath } from "obsidian";

import { SEARCH_CONSTANTS, REGEX_CONSTANTS } from "../constants";
import { GraphService } from "../services/GraphService";
import { VaultIntelligenceSettings, DEFAULT_SETTINGS, DEFAULT_SYSTEM_PROMPT } from "../settings";
import { FileTools } from "../tools/FileTools";
import { ToolRegistry } from "../tools/ToolRegistry";
import { IEmbeddingClient, IToolDefinition, StreamChunk, ToolCall, ToolResult, UnifiedMessage, IReasoningClient, IModelProvider } from "../types/providers";
import { VaultSearchResult } from "../types/search";
import { logger } from "../utils/logger";
import { ContextAssembler } from "./ContextAssembler";
import { McpClientManager } from "./McpClientManager";
import { ModelRegistry } from "./ModelRegistry";
import { ProviderRegistry } from "./ProviderRegistry";
import { SearchOrchestrator } from "./SearchOrchestrator";

export interface ChatMessage {
    contextFiles?: string[];
    createdFiles?: string[];
    error?: string;
    isCancelled?: boolean;
    rawContent?: unknown[];
    role: "user" | "model" | "system" | "tool";
    spotlightResults?: VaultSearchResult[];
    text: string;
    thought?: string;
    toolCalls?: ToolCall[];
    toolResults?: ToolResult[];
}

// Internal interface for accessing non-public Obsidian API properties
interface InternalWorkspaceLeaf {
    parent?: {
        activeLeaf?: WorkspaceLeaf;
        type: string;
    };
}

/**
 * Service responsible for orchestrating the AI agent activities.
 * It manages tool execution, chat history, and context assembly.
 */
export class AgentService {
    private app: App;
    private embeddingService: IEmbeddingClient;
    private graphService: GraphService;
    private providerRegistry: ProviderRegistry;
    private settings: VaultIntelligenceSettings;

    private contextAssembler: ContextAssembler;
    private searchOrchestrator: SearchOrchestrator;
    private toolRegistry: ToolRegistry;

    constructor(
        app: App,
        providerRegistry: ProviderRegistry,
        graphService: GraphService,
        embeddingService: IEmbeddingClient,
        settings: VaultIntelligenceSettings,
        mcpClientManager: McpClientManager
    ) {
        this.app = app;
        this.providerRegistry = providerRegistry;
        this.graphService = graphService;
        this.embeddingService = embeddingService;
        this.settings = settings;

        // Initialize with default providers from registry
        const initialClient = this.providerRegistry.getReasoningClient();
        const initialProvider = this.providerRegistry.getModelProvider();

        this.searchOrchestrator = new SearchOrchestrator(app, graphService, initialClient, this.embeddingService, settings);
        this.contextAssembler = new ContextAssembler(app, graphService, settings, mcpClientManager);

        const fileTools = new FileTools(app);
        this.toolRegistry = new ToolRegistry(
            app,
            settings,
            initialClient,
            initialProvider,
            graphService,
            this.searchOrchestrator,
            this.contextAssembler,
            fileTools,
            mcpClientManager
        );
    }

    public getSearchOrchestrator(): SearchOrchestrator {
        return this.searchOrchestrator;
    }

    public async reflexSearch(query: string, limit: number): Promise<VaultSearchResult[]> {
        return this.searchOrchestrator.searchReflex(query, limit);
    }

    public async chat(
        messages: ChatMessage[],
        currentPrompt: string,
        contextFiles: TFile[] = [],
        options: { modelId?: string; enableCodeExecution?: boolean; enableUrlContext?: boolean; enableWebSearch?: boolean; enableAgentWriteAccess?: boolean; signal?: AbortSignal } = {}
    ): Promise<{ createdFiles: string[]; files: string[]; text: string }> {
        const stream = this.chatStream(messages, currentPrompt, contextFiles, options);
        let finalText = "";
        let finalFiles: string[] = [];
        let finalCreatedFiles: string[] = [];

        for await (const chunk of stream) {
            if (chunk.replaceText !== undefined) finalText = chunk.replaceText;
            else if (chunk.text) finalText += chunk.text;
            if (chunk.isDone) {
                finalFiles = chunk.files || [];
                finalCreatedFiles = chunk.createdFiles || [];
            }
        }
        return { createdFiles: finalCreatedFiles, files: finalFiles, text: finalText };
    }

    public async *chatStream(
        messages: ChatMessage[],
        currentPrompt: string,
        contextFiles: TFile[] = [],
        options: { modelId?: string; enableCodeExecution?: boolean; enableUrlContext?: boolean; enableWebSearch?: boolean; enableAgentWriteAccess?: boolean; signal?: AbortSignal } = {}
    ): AsyncIterableIterator<StreamChunk> {

        const reasoningClient: IReasoningClient = this.providerRegistry.getReasoningClient(options.modelId);
        const provider: IModelProvider = this.providerRegistry.getModelProvider(options.modelId);

        // Update ToolRegistry for this turn if model changed
        this.toolRegistry.updateProvider(reasoningClient, provider);

        const activeFile = this.app.workspace.getActiveFile();
        const explicitlyMentionedCount = contextFiles.length;

        // ALWAYS inject the active file first so the LLM knows what the "current document" is.
        if (activeFile && !contextFiles.some(f => f.path === activeFile.path)) {
            contextFiles.unshift(activeFile);
        }

        // Only auto-inject other visible leaves if NO files were explicitly mentioned via `@`
        if (explicitlyMentionedCount === 0) {
            this.app.workspace.iterateRootLeaves((leaf) => {
                const view = leaf.view;
                if (view instanceof MarkdownView) {
                    const file = view.file;
                    if (file && (!activeFile || file.path !== activeFile.path)) {
                        const internalLeaf = leaf as unknown as InternalWorkspaceLeaf;
                        const isVisible = internalLeaf.parent?.type === "tabs" ? internalLeaf.parent.activeLeaf === leaf : true;

                        if (isVisible && !contextFiles.some(f => f.path === file.path)) {
                            contextFiles.push(file);
                        }
                    }
                }
            });
        }

        if (contextFiles.length > 0) {
            logger.info(`[Agent] Context assembled. Total files: ${contextFiles.length}. (Active: ${activeFile?.path ?? "None"}).`);
        }

        const usedFiles = new Set<string>();
        const createdFiles = new Set<string>();
        const wereFilesExplicitlyMentioned = explicitlyMentionedCount > 0;

        // 1. Filter out UI-only messages (e.g., Spotlight results) and map to UnifiedMessage
        const filteredMessages: UnifiedMessage[] = messages
            .filter(msg => {
                // Keep if it has text, tool calls, or tool results
                if (msg.text?.trim() || msg.toolCalls?.length || msg.toolResults?.length) return true;
                // Keep if it has raw content parts
                if (msg.rawContent?.length) return true;
                return false;
            })
            .map(h => ({
                content: h.text,
                rawContent: h.rawContent,
                role: h.role,
                toolCalls: h.toolCalls,
                toolResults: h.toolResults
            }));

        // 2. De-duplicate if the UI already added this message to history
        // We look for any 'user' role message that matches the current prompt exactly.
        const duplicateIndex = filteredMessages.findIndex(m => m.role === 'user' && m.content === currentPrompt);
        if (duplicateIndex !== -1) {
            filteredMessages.splice(duplicateIndex, 1);
        }

        let formattedHistory: UnifiedMessage[] = filteredMessages;

        if (contextFiles.length > 0) {
            const fileResults: VaultSearchResult[] = contextFiles.map(f => ({
                isKeywordMatch: true,
                path: f.path,
                score: 1.0
            }));

            const activeModel = options.modelId || this.settings.chatModel;
            const totalTokens = ModelRegistry.resolveContextBudget(activeModel, this.settings.modelContextOverrides, this.settings.contextWindowTokens);
            const contextBudget = Math.floor(totalTokens * SEARCH_CONSTANTS.CONTEXT_SAFETY_MARGIN);

            const { context, usedFiles: assembledFiles } = await this.contextAssembler.assemble(fileResults, currentPrompt, contextBudget);

            if (context) {
                if (wereFilesExplicitlyMentioned && assembledFiles.length < contextFiles.length) {
                    const skipped = contextFiles.length - assembledFiles.length;
                    currentPrompt = `[SYSTEM NOTE: ${skipped} context files were skipped because they exceeded the context window budget for the selected model.]\n\n` + currentPrompt;
                }
                
                const activeFileNote = activeFile ? `\n\n[SYSTEM NOTE: The user's CURRENTLY OPEN and ACTIVE document is: "${activeFile.path}". If the user asks to "update the current document", "fix the current note", or "add to this", they are referring EXACTLY to "${activeFile.path}".]` : "";

                currentPrompt = `The following notes were automatically injected from your workspace context:\n${context}${activeFileNote}\n\nUser Query: ${currentPrompt}`;

                // ONLY track files that actually survived the budget
                assembledFiles.forEach(f => usedFiles.add(normalizePath(f)));
            }
        }

        const currentDate = new Date().toDateString();
        const rawSystemInstruction = this.settings.systemInstruction ?? DEFAULT_SYSTEM_PROMPT;
        let systemInstruction = (rawSystemInstruction || "").replace("{{DATE}}", currentDate);
        systemInstruction = systemInstruction.replace("{{LANGUAGE}}", this.settings.agentLanguage || "English (US)");

        const tools: IToolDefinition[] = await this.toolRegistry.getTools({
            enableCodeExecution: options.enableCodeExecution,
            enableWebSearch: options.enableWebSearch
        });
        const validToolNames = new Set(tools.map(t => t.name));

        const activeModelStr = options.modelId || this.settings.chatModel;
        const providerName = (activeModelStr || "").startsWith("ollama/") || (activeModelStr || "").startsWith("local/") ? "ollama" : "gemini";
        const supportsWeb = provider.supportsWebGrounding;

        // Dynamic System Prompt Variables safely injected via templates rather than fragile regex
        if (supportsWeb && tools.find(t => t.name === "google_search")) {
            systemInstruction = systemInstruction.replace("{{VERIFICATION_RULES}}", "2. **Verification**: When users ask for facts, ALWAYS verify them against real-world data using 'google_search' unless explicitly told to rely only on notes.");
        } else {
            systemInstruction = systemInstruction.replace("{{VERIFICATION_RULES}}", "2. **Verification**: Always double-check facts against the user's notes.");
            systemInstruction = systemInstruction.replace(/\s*-\s*Use\s*'?google_search'?[\s\S]*?(?=\n)/gi, "");
            logger.debug("[Agent] Dynamically stripped google_search instruction from system prompt", { provider: providerName });
        }

        // Filter out "ghost tools" from history (e.g. disabled MCP servers) to prevent LLM hallucinations
        formattedHistory = formattedHistory.map(msg => {
            const newMsg = { ...msg };
            if (newMsg.toolCalls) {
                newMsg.toolCalls = newMsg.toolCalls.filter(tc => validToolNames.has(tc.name));
                if (newMsg.toolCalls.length === 0) newMsg.toolCalls = undefined;
            }
            if (newMsg.toolResults) {
                newMsg.toolResults = newMsg.toolResults.filter(tr => validToolNames.has(tr.name));
                if (newMsg.toolResults.length === 0) newMsg.toolResults = undefined;
            }
            return newMsg;
        }).filter(msg => msg.content?.trim() || msg.toolCalls?.length || msg.toolResults?.length || msg.rawContent?.length);

        formattedHistory.push({ content: currentPrompt, role: "user" });
        logger.debug(`[Agent] History turns: ${formattedHistory.length}`);
        const debugPrompt = currentPrompt.length > 500 ? currentPrompt.substring(0, 500) + "..." : currentPrompt;
        logger.debug(`[Agent] Final Enriched Prompt Preview:\n${debugPrompt}`);

        try {
            const maxLoops = this.settings?.maxAgentSteps ?? DEFAULT_SETTINGS.maxAgentSteps;
            yield* this.processToolLoop(
                formattedHistory,
                systemInstruction,
                tools,
                createdFiles,
                usedFiles,
                options,
                maxLoops
            );

            if (options.signal?.aborted) {
                yield { error: "Agent explicitly stopped.", isCancelled: true, isDone: true };
            } else {
                yield {
                    createdFiles: Array.from(createdFiles),
                    files: Array.from(usedFiles),
                    isDone: true
                };
            }
        } catch (e: unknown) {
            logger.error("Error in AgentService chatStream", e);
            const errorMessage = e instanceof Error ? e.message : String(e);

            if (errorMessage.includes("400") || errorMessage.includes("API key")) {
                yield { error: `I encountered an error connecting to the AI provider (Status 400). Please check your API key. Details: ${errorMessage}`, isDone: true };
            } else if ((e instanceof DOMException && e.name === "AbortError") || errorMessage.includes("AbortError") || errorMessage.includes("aborted")) {
                yield { error: "Agent explicitly stopped.", isCancelled: true, isDone: true };
            } else {
                yield { error: errorMessage, isDone: true };
            }
        }
    }

    private async *processToolLoop(
        formattedHistory: UnifiedMessage[],
        systemInstruction: string,
        tools: IToolDefinition[],
        createdFiles: Set<string>,
        usedFiles: Set<string>,
        options: { modelId?: string; enableCodeExecution?: boolean; enableUrlContext?: boolean; enableWebSearch?: boolean; enableAgentWriteAccess?: boolean; signal?: AbortSignal },
        maxLoops: number
    ): AsyncIterableIterator<StreamChunk> {
        const reasoningClient = this.providerRegistry.getReasoningClient(options.modelId);
        let loops = 0;

        while (loops < maxLoops) {
            if (options.signal?.aborted) break;

            let modelResponseRawContent: unknown[] | undefined;

            const stream = reasoningClient.generateMessageStream(formattedHistory, {
                enableCodeExecution: options.enableCodeExecution,
                enableUrlContext: options.enableUrlContext,
                enableWebSearch: options.enableWebSearch,
                modelId: options.modelId,
                signal: options.signal,
                systemInstruction: systemInstruction,
                tools: tools.length > 0 ? tools : undefined
            });

            let loopText = "";
            const loopToolCalls: ToolCall[] = [];

            for await (const chunk of stream) {
                if (options.signal?.aborted) break;

                if (chunk.toolCalls) {
                    loopToolCalls.push(...chunk.toolCalls);
                }
                if (chunk.replaceText !== undefined) {
                    loopText = chunk.replaceText;
                    yield { replaceText: chunk.replaceText };
                }
                if (chunk.text !== undefined) {
                    loopText += chunk.text;
                    yield { text: chunk.text };
                }
                if (chunk.rawContent) {
                    modelResponseRawContent = chunk.rawContent;
                }
            }

            if (options.signal?.aborted) break;

            const modelResponse: UnifiedMessage = {
                content: loopText,
                rawContent: modelResponseRawContent,
                role: "model",
                toolCalls: loopToolCalls.length > 0 ? loopToolCalls : undefined
            };
            formattedHistory.push(modelResponse);

            // Yield the final model message metadata so UI can store it for follow-up turns
            yield { 
                rawContent: modelResponseRawContent,
                toolCalls: loopToolCalls.length > 0 ? loopToolCalls : undefined
            };

            if (loopToolCalls.length > 0) {
                const toolStatus = `Thinking: Using tools (${loopToolCalls.map(c => c.name).join(", ")})...`;
                yield { status: toolStatus };

                const toolPromises = loopToolCalls.map(async (call) => {
                    try {
                        const args = call.args || {};
                        const functionResponse = await this.toolRegistry.execute({
                            args: args,
                            createdFiles: createdFiles,
                            enableAgentWriteAccess: options.enableAgentWriteAccess,
                            enableCodeExecution: options.enableCodeExecution,
                            modelId: options.modelId,
                            name: call.name,
                            signal: options.signal,
                            usedFiles: usedFiles
                        });

                        return {
                            id: call.id,
                            name: call.name,
                            result: functionResponse,
                            thought_signature: call.thought_signature
                        };
                    } catch (error) {
                        logger.error(`[Agent] Tool execution failed for ${call.name}`, error);
                        return {
                            id: call.id,
                            name: call.name,
                            result: { error: String(error) },
                            thought_signature: call.thought_signature
                        };
                    }
                });

                const completedParts = await Promise.all(toolPromises);

                if (completedParts.length > 0) {
                    const toolMsg: UnifiedMessage = {
                        content: "",
                        role: "tool",
                        toolResults: completedParts.map(p => ({
                            id: p.id,
                            name: p.name,
                            result: p.result,
                            thought_signature: p.thought_signature
                        }))
                    };
                    formattedHistory.push(toolMsg);
                    
                    // Yield tool results to UI for history preservation
                    yield { toolResults: toolMsg.toolResults };
                } else {
                    break;
                }
            } else {
                break;
            }
            loops++;
        }

        if (loops >= maxLoops) {
            yield { text: "\n\nI'm sorry, I reached the step limit before finding a definitive answer. Please try rephrasing or check your settings." };
        }
    }

    /**
     * Prepares the context for a chat message by resolving file references (`@Filename`).
     * This separates the view logic from the business logic.
     * @returns Object containing resolved context files, the clean message, and any warnings.
     */
    public async prepareContext(inputMessage: string, modelId?: string): Promise<{ contextFiles: TFile[], cleanMessage: string, warnings: string[] }> {
        const resultFiles: TFile[] = [];
        const warnings: string[] = [];
        let message = inputMessage;

        // Regex for @[["Filename"]] or @[[Filename|Alias]] or @Filename 
        // Updated regex to be more lenient with characters and handle quoted strings better
        const mentionRegex = REGEX_CONSTANTS.MENTION;
        const matches = Array.from(inputMessage.matchAll(mentionRegex));

        // We will collect all potential files
        const allPotentialFiles: TFile[] = [];

        // We don't remove the mentions from the message, as the model might need them to understand "Look at @File"
        // But we could optionally clean them. ResearchChatView didn't clean them, so we won't either.

        for (const match of matches) {
            const pathOrName = match[1] || match[2] || match[3];
            if (pathOrName) {
                // First try direct path
                let abstractFile = this.app.vault.getAbstractFileByPath(pathOrName);

                // If not found, try resolving as a link (for files)
                if (!abstractFile) {
                    abstractFile = this.app.metadataCache.getFirstLinkpathDest(pathOrName, "");
                }

                if (abstractFile instanceof TFile) {
                    if (!allPotentialFiles.some(f => f.path === abstractFile.path)) {
                        allPotentialFiles.push(abstractFile);
                    }
                } else if (abstractFile instanceof TFolder) {
                    // Expand folder into files recursively
                    await this.processFolderContext(abstractFile, message, allPotentialFiles, warnings, modelId);
                }
            }
        }

        resultFiles.push(...allPotentialFiles);

        return { cleanMessage: message, contextFiles: resultFiles, warnings };
    }

    private async processFolderContext(folder: TFolder, query: string, collector: TFile[], warnings: string[], modelId?: string) {
        // TFolder is a type from 'obsidian', assuming it's imported at the top of the file.

        const folderPath = folder.path;
        const folderFiles = this.app.vault.getMarkdownFiles().filter(f =>
            f.path.startsWith(folderPath + "/") || f.path === folderPath
        );

        // Try similarity search within the folder first
        const folderPaths = folderFiles.map(f => f.path);
        const similarityResults = await this.graphService.searchInPaths(query, folderPaths, 100);

        let sortedFiles: TFile[];
        if (similarityResults.length > 0) {
            const resultPathMap = new Map(similarityResults.map((r, i) => [r.path, i]));
            const matchedFiles = folderFiles.filter(f => resultPathMap.has(f.path));
            matchedFiles.sort((a, b) => (resultPathMap.get(a.path) ?? 0) - (resultPathMap.get(b.path) ?? 0));

            // Include unmatched files (not indexed yet) by recency at the end
            const unmatchedFiles = folderFiles.filter(f => !resultPathMap.has(f.path));
            unmatchedFiles.sort((a, b) => b.stat.mtime - a.stat.mtime);

            sortedFiles = [...matchedFiles, ...unmatchedFiles];
        } else {
            // Fallback to recency if no similarity results (e.g. index not ready)
            sortedFiles = [...folderFiles];
            sortedFiles.sort((a, b) => b.stat.mtime - a.stat.mtime);
        }

        // Calculate budget
        // We allocate 50% of the total context window for explicit folder mentions to leave room for history/responses
        const activeModel = modelId || this.settings.chatModel;
        const totalTokens = ModelRegistry.resolveContextBudget(activeModel, this.settings.modelContextOverrides, this.settings.contextWindowTokens);
        const budgetTokens = Math.floor(totalTokens * SEARCH_CONSTANTS.CONTEXT_SAFETY_MARGIN * 0.5); // 50% for folder context
        const budgetChars = budgetTokens * SEARCH_CONSTANTS.CHARS_PER_TOKEN_ESTIMATE;

        // Starvation Protection
        // If only one document, let it take the whole budget. If multiple, apply soft limit.
        const effectiveLimitRatio = sortedFiles.length <= 1 ? 1.0 : SEARCH_CONSTANTS.SINGLE_DOC_SOFT_LIMIT_RATIO;
        const singleDocSoftLimitTokens = Math.floor(budgetTokens * effectiveLimitRatio);
        const singleDocSoftLimitChars = Math.floor(budgetChars * effectiveLimitRatio);

        logger.debug(`[ContextAssembler] Folder "${folder.name}" Budget: ${budgetTokens} tokens. Soft Cap (${(effectiveLimitRatio * 100).toFixed(0)}%): ${singleDocSoftLimitTokens} tokens.`);

        const maxFiles = this.settings.contextMaxFiles || SEARCH_CONSTANTS.DEFAULT_CONTEXT_MAX_FILES;

        let currentSize = 0;
        const filesToInclude: TFile[] = [];

        for (const file of sortedFiles) {
            const size = file.stat.size;
            // Apply soft limit for individual files if there are multiple files
            if (sortedFiles.length > 1 && currentSize + size > singleDocSoftLimitChars && filesToInclude.length > 0) {
                logger.debug(`[ContextAssembler] Skipping "${file.name}" due to single document soft limit. Current size: ${currentSize}, File size: ${size}, Soft limit: ${singleDocSoftLimitChars}`);
                continue; // Skip this file but continue trying others
            }
            if (currentSize + size > budgetChars || filesToInclude.length >= maxFiles) {
                logger.debug(`[ContextAssembler] Stopping file inclusion for folder "${folder.name}". Current size: ${currentSize}, File size: ${size}, Budget: ${budgetChars}, Files included: ${filesToInclude.length}, Max files: ${maxFiles}`);
                break;
            }

            filesToInclude.push(file);
            currentSize += size;
        }

        if (filesToInclude.length < folderFiles.length) {
            const method = similarityResults.length > 0 ? "similarity-ranked" : "most recent";
            const limitReason = filesToInclude.length >= maxFiles ? `Max file limit (${maxFiles})` : `Context limit (~${totalTokens.toLocaleString()} tokens)`;
            warnings.push(`${limitReason} reached for folder "${folder.name}". Included ${filesToInclude.length} ${method} files.`);
        }

        filesToInclude.forEach(f => {
            if (!collector.some(existing => existing.path === f.path)) {
                collector.push(f);
            }
        });
    }
}
