import { App, normalizePath } from "obsidian";
import { z } from "zod";

import { SEARCH_CONSTANTS } from "../constants";
import { VaultIntelligenceSettings } from "../settings/types";
import { IEmbeddingClient, IReasoningClient } from "../types/providers";
import { VaultSearchResult } from "../types/search";
import { logger } from "../utils/logger";
import { GraphService } from "./GraphService";
import { ScoringStrategy } from "./ScoringStrategy";

/**
 * Service that orchestrates hybrid search across the vault.
 * Combines vector search results with keyword matching for optimal recall and precision.
 */
export class SearchOrchestrator {
    private app: App;
    private graphService: GraphService;
    private reasoningClient: IReasoningClient;
    private embeddingService: IEmbeddingClient;
    private settings: VaultIntelligenceSettings;
    private scoringStrategy: ScoringStrategy;

    // FIX: Ensure constructor accepts 5 arguments including embeddingService
    constructor(
        app: App,
        graphService: GraphService,
        reasoningClient: IReasoningClient,
        embeddingService: IEmbeddingClient,
        settings: VaultIntelligenceSettings
    ) {
        this.app = app;
        this.graphService = graphService;
        this.reasoningClient = reasoningClient;
        this.embeddingService = embeddingService;
        this.settings = settings;
        this.scoringStrategy = new ScoringStrategy();
    }

    /**
     * Updates the reasoning client used for Analyst search (Loop 2).
     */
    public updateReasoningClient(client: IReasoningClient): void {
        this.reasoningClient = client;
    }

    /**
     * Performs a hybrid search (Vector + Keyword) on the vault.
     * @param query - The user's search query.
     * @param limit - Maximum number of final results to return.
     * @param options - Search options (e.g. deep: false to skip AI reranking).
     * @returns A promise resolving to a ranked list of search results.
     */
    public async search(query: string, limit: number, options: { deep?: boolean; signal?: AbortSignal } = {}): Promise<VaultSearchResult[]> {
        if (options.signal?.aborted) throw new DOMException("Aborted", "AbortError");
        // DUAL-LOOP LOGIC:
        // If Deep search is requested (or default is enabled), we try the Analyst (Loop 2) first.
        // It provides deeper, graph-expanded, and AI-ranked results.
        const deep = options.deep ?? this.settings.enableDualLoop;

        if (deep) {
            try {
                const analystResults = await this.searchAnalyst(query, options.signal);
                if (analystResults.length > 0) {
                    return analystResults.slice(0, limit);
                }
                logger.warn("[SearchOrchestrator] Analyst loop returned no results. Falling back to Reflex.");
            } catch (error) {
                logger.error("[SearchOrchestrator] Analyst loop failed. Falling back to Reflex.", error);
            }
        }

        // Fallback or Default: Reflex Loop (Loop 1)
        return this.searchReflex(query, limit, options.signal);
    }

    /**
     * DUAL-LOOP: Loop 1 (Reflex)
     * Fast, local, parameter-free search.
     */
    public async searchReflex(query: string, limit: number, signal?: AbortSignal): Promise<VaultSearchResult[]> {
        if (signal?.aborted) throw new DOMException("Aborted", "AbortError");
        if (!query || query.trim().length === 0) {
            return [];
        }

        logger.info(`[SearchOrchestrator] Reflex Search: "${query}"`);

        const workerLimit = Math.max(limit * 2, 50);
        const vectorResults = await this.graphService.search(query, workerLimit);
        if (signal?.aborted) throw new DOMException("Aborted", "AbortError");
        
        const keywordResults = await this.graphService.keywordSearch(query, workerLimit);
        if (signal?.aborted) throw new DOMException("Aborted", "AbortError");

        const vResults: VaultSearchResult[] = vectorResults.map(r => ({
            excerpt: r.excerpt,
            isKeywordMatch: false,
            isTitleMatch: false,
            path: r.path,
            score: r.score,
            tokenCount: r.tokenCount
        }));

        const kResults: VaultSearchResult[] = keywordResults.map(r => ({
            excerpt: r.excerpt,
            isKeywordMatch: true,
            isTitleMatch: false,
            path: r.path,
            score: r.score,
            tokenCount: r.tokenCount
        }));

        const merged = this.mergeAndRank(vResults, kResults, limit);
        return merged;
    }

    /**
     * DUAL-LOOP: Loop 2 (Analyst)
     * Deep, AI-driven re-ranking.
     */
    public async searchAnalyst(query: string, signal?: AbortSignal): Promise<VaultSearchResult[]> {
        if (signal?.aborted) throw new DOMException("Aborted", "AbortError");
        logger.info(`[SearchOrchestrator] Analyst Search: "${query}"`);

        if (this.settings.enableDualLoop) {
            // Dual-Loop: Reflex (Loop 1) is handled by UI. This is Analyst (Loop 2).
            const { vector: queryVector } = await this.embeddingService.embedQuery(query);
            if (signal?.aborted) throw new DOMException("Aborted", "AbortError");

            // 1. Build Payload (Graph + Vector + Keyword)
            const payload = await this.graphService.buildPriorityPayload(queryVector, query);

            // 2. Re-Rank using generic structured output
            const ReRankSchema = z.array(
                z.object({
                    id: z.string().describe("The path of the document"),
                    reasoning: z.string().describe("Why this document is relevant"),
                    score: z.number().describe("Relevance score between 0 and 1"),
                    tokenCount: z.number().optional()
                })
            );

            const instruction = `You are a search analyst. Given a search query and a list of potentially relevant documents, score and re-rank them based on absolute relevance. Document Context:\n${JSON.stringify(payload)}`;
            
            const reranked = await this.reasoningClient.generateStructured(
                [{ content: query, role: "user" }],
                ReRankSchema,
                {
                    jsonSchema: {
                        items: {
                            properties: {
                                id: { description: "The path of the document", type: "string" },
                                reasoning: { description: "Why this document is relevant", type: "string" },
                                score: { description: "Relevance score between 0 and 1", type: "number" },
                                tokenCount: { type: "number" }
                            },
                            required: ["id", "reasoning", "score"],
                            type: "object"
                        },
                        type: "array"
                    },
                    modelId: this.settings.reRankingModel,
                    signal: signal,
                    systemInstruction: instruction
                }
            );

            if (signal?.aborted) throw new DOMException("Aborted", "AbortError");

            // 3. Map to VaultSearchResult
            return reranked.map(r => ({
                content: r.reasoning,
                path: normalizePath(String(r.id).split('#')[0] || ""), // Normalize AI-provided path
                score: r.score,
                tokenCount: r.tokenCount
            } as VaultSearchResult));
        }
        return []; // Return empty if dual loop is not enabled
    }

    /**
     * Merges vector-based results with keyword-based results, applying boosts where they overlap.
     * @param vectorResults - Results from the vector search.
     * @param keywordResults - Results from the keyword search.
     * @param limit - Maximum number of results to return after merging.
     * @returns Array of merged and ranked search results.
     * @private
     */
    private mergeAndRank(vectorResults: VaultSearchResult[], keywordResults: VaultSearchResult[], limit: number): VaultSearchResult[] {
        const mergedMap = new Map<string, VaultSearchResult>();

        // Add Vector Results
        for (const res of vectorResults) {
            mergedMap.set(res.path, res);
        }

        // We use a Calibration Factor to map unbounded BM25 scores into 0-1 range.
        // A sigmoid-like function: score / (score + K) allows high scores to approach 1.0 
        // without capping, preserving ranking granularity.
        const K = this.settings.keywordWeight;

        for (const res of keywordResults) {
            const existing = mergedMap.get(res.path);
            const normalizedKeywordScore = res.score / (res.score + K);

            if (existing !== undefined) {
                // Blend scores. We give vector results slightly more weight by dividing by HYBRID_VECTOR_DIVIDER.
                existing.score = (existing.score + normalizedKeywordScore) / SEARCH_CONSTANTS.HYBRID_VECTOR_DIVIDER;
                existing.isKeywordMatch = true;
            } else {
                mergedMap.set(res.path, { ...res, score: normalizedKeywordScore });
            }
        }

        const minScore = this.settings.minSimilarityScore;

        const finalResults = Array.from(mergedMap.values())
            .filter(r => r.score >= minScore)
            .sort((a, b) => b.score - a.score)
            .slice(0, limit);

        return finalResults;
    }
}