import type { RunTree } from "langsmith";
import { executeCompleteText } from "./inference/execute";
import { getDefaultModelForFeature } from "./inference/config";
import type { UserApiKeys } from "./llm/types";
import { WEB_SEARCH_ROUTING_SYSTEM_PROMPT } from "../prompts/chat/routing/web-search-routing";

const WEB_SEARCH_INTENT_RE =
    /\b(search|browse|google|look\s+up|web|internet|online|latest|recently|recent|current|today|news|source(?:s)?|cite(?:s|d)?|verify|fact[- ]?check)\b/i;

/** Person/entity + said/talked phrasing usually needs live news context. */
const RECENT_STATEMENT_RE =
    /\b(what did|what has|what does)\b.{0,80}\b(say|said|talk|talked|speak|spoke|mention|mentioned|comment|commented)\b/i;

const REALTIME_FACT_RE =
    /\b(price|prices|stock price|share price|exchange rate|market cap|trading at|trading on|how much(?:\s+is)?|worth|value of)\b/i;

const REALTIME_MARKET_RE =
    /\bhow much\b.{0,80}\b(stock|stocks|shares?|equity|crypto|cryptocurrency|coin|coins)\b/i;

const MARKET_ASSET_RE =
    /\b(stock|stocks|shares?|equity|crypto|cryptocurrency|coin|coins?|bitcoin|btc|ethereum|eth|solana|sol|tao|exchange rate|market cap|trading at|trading on)\b/i;

/** Ambiguous definitions in a geo/news context (e.g. trending acronyms like CJP in India). */
const AMBIGUOUS_GEO_DEFINITION_RE =
    /\b(what is|what's|who is|who's|what does|what do)\b.{0,80}\b(in|for)\b.{0,40}\b(india|usa|u\.?s\.?a?|uk|china|europe|world)\b/i;

const CURRENT_LEGAL_RE =
    /\b(current|latest|recent|updated|in force|as of today|as of now)\b.{0,80}\b(law|regulation|statute|case law|guidance|rule|deadline|filing|fee|authority|ministry|court|article)\b/i;

/** Questions that may need live context when heuristics do not already match. */
const ROUTING_CANDIDATE_RE =
    /\b(what is|what's|who is|who's|what does|what do|meaning of|stand for|stands for|trending|why is|why are|explain|define|tell me about|what happened|what's happening|what is happening)\b/i;

const LOCAL_FAST_PATH_RE =
    /^\s*(hi|hello|hey|thanks|thank you|ok|okay|yes|no)\s*[!.?]*\s*$/i;

const DOCUMENT_SCOPE_RE =
    /\b(this|the|above|attached|uploaded|following)\b.{0,40}\b(document|docx|file|clause|section|paragraph|contract|agreement|draft|exhibit|schedule)\b/i;

export type WebSearchRoutingSource = "heuristic" | "llm" | "skip";

export type WebSearchRoutingResult = {
    needsWebSearch: boolean;
    source: WebSearchRoutingSource;
    llmRaw?: string;
};

export function detectWebSearchIntentHeuristic(
    lastUserText: string,
    options: { hasAttachedDocs?: boolean } = {},
): boolean {
    const hasAttachedDocs = options.hasAttachedDocs ?? false;
    const realtimeFact = REALTIME_FACT_RE.test(lastUserText);
    const realtimeMarket = REALTIME_MARKET_RE.test(lastUserText);

    if (
        hasAttachedDocs &&
        realtimeFact &&
        !realtimeMarket &&
        !WEB_SEARCH_INTENT_RE.test(lastUserText) &&
        !CURRENT_LEGAL_RE.test(lastUserText) &&
        !MARKET_ASSET_RE.test(lastUserText)
    ) {
        return false;
    }

    return (
        WEB_SEARCH_INTENT_RE.test(lastUserText) ||
        RECENT_STATEMENT_RE.test(lastUserText) ||
        CURRENT_LEGAL_RE.test(lastUserText) ||
        realtimeFact ||
        realtimeMarket ||
        AMBIGUOUS_GEO_DEFINITION_RE.test(lastUserText)
    );
}

function shouldSkipLlmRouting(lastUserText: string, hasAttachedDocs: boolean): boolean {
    const trimmed = lastUserText.trim();
    if (!trimmed) return true;
    if (trimmed.length < 8) return true;
    if (LOCAL_FAST_PATH_RE.test(trimmed)) return true;
    if (hasAttachedDocs && DOCUMENT_SCOPE_RE.test(trimmed)) return true;
    if (!ROUTING_CANDIDATE_RE.test(trimmed)) return true;
    return false;
}

function parseRouterResponse(raw: string): boolean {
    const token = raw.trim().toUpperCase().split(/\s+/)[0] ?? "";
    if (token.startsWith("WEB")) return true;
    if (token.startsWith("LOCAL")) return false;
    return /\bWEB\b/.test(raw.toUpperCase());
}

export async function classifyWebSearchNeed(params: {
    userText: string;
    apiKeys?: UserApiKeys;
    trace?: RunTree | null;
    hasAttachedDocs?: boolean;
}): Promise<WebSearchRoutingResult> {
    const userText = params.userText.trim();
    const hasAttachedDocs = params.hasAttachedDocs ?? false;

    if (hasAttachedDocs) {
        return { needsWebSearch: false, source: "skip" };
    }

    if (detectWebSearchIntentHeuristic(userText, { hasAttachedDocs })) {
        return { needsWebSearch: true, source: "heuristic" };
    }

    if (shouldSkipLlmRouting(userText, hasAttachedDocs)) {
        return { needsWebSearch: false, source: "skip" };
    }

    if (process.env.WEB_SEARCH_LLM_ROUTING === "0") {
        return { needsWebSearch: false, source: "skip" };
    }

    try {
        const callRouter = async (maxTokens: number) =>
            executeCompleteText({
                model: getDefaultModelForFeature("web-search-routing"),
                featureId: "web-search-routing",
                systemPrompt: WEB_SEARCH_ROUTING_SYSTEM_PROMPT,
                user: userText,
                maxTokens,
                temperature: 0,
                lightweight: true,
                apiKeys: params.apiKeys,
                trace: params.trace,
            });

        // Proxy models need a non-trivial max_tokens budget; very small limits return null content.
        let raw = await callRouter(128);
        if (!raw.trim()) {
            raw = await callRouter(256);
        }

        const needsWebSearch = raw.trim()
            ? parseRouterResponse(raw)
            : true;
        return { needsWebSearch, source: "llm", llmRaw: raw };
    } catch {
        return { needsWebSearch: false, source: "skip" };
    }
}