import { test, describe, it } from "node:test";
import assert from "node:assert";
import {
    parseRequestedAssistantMode,
    resolveTurnMode,
    resolveTurnModeAsync,
    applySaudiLawRouting,
    type RequestedAssistantMode,
} from "../../src/lib/documentMode";
import { buildSaudiLawRoutingResult } from "../../src/lib/saudiLawRouting";
import { isToolAllowedForMode } from "../../src/lib/chatTools";

describe("documentModeConfirmation", () => {
    it("validates every explicit requested mode and rejects unknown values", () => {
        const modes: RequestedAssistantMode[] = [
            "auto",
            "answer",
            "comments",
            "document_edit",
            "workflow_review",
            "document_generation",
            "web_search",
        ];
        for (const mode of modes) {
            assert.strictEqual(parseRequestedAssistantMode(mode), mode);
        }
        assert.throws(
            () => parseRequestedAssistantMode("invalid"),
            /invalid_requested_mode/,
        );
    });

    it("returns explicit mode directly", () => {
        const decision = resolveTurnMode({
            requestedMode: "comments",
            messages: [],
            docIndex: {},
        });
        assert.strictEqual(decision.turnMode, "comments");
        assert.strictEqual(decision.requestedMode, "comments");
        assert.strictEqual(decision.requiresConfirmation, false);
    });

    it("requires confirmation for inferred direct_docx_mutation", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [{ role: "user", content: "rewrite this section to be better" }],
            requestedTrackChanges: true,
            docIndex: {
                "doc-0": { document_id: "doc-0", filename: "foo.docx", file_type: "docx" }
            },
        });
        assert.strictEqual(decision.turnMode, "document_direct_edit");
        assert.strictEqual(decision.source, "inferred");
        assert.strictEqual(decision.requiresConfirmation, true);
    });

    it("requires confirmation for inferred docx_mutation_with_web_search", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [{ role: "user", content: "rewrite this section based on latest law" }],
            requestedTrackChanges: true,
            docIndex: {
                "doc-0": { document_id: "doc-0", filename: "foo.docx", file_type: "docx" }
            },
        });
        assert.strictEqual(decision.turnMode, "document_direct_edit");
        assert.strictEqual(decision.source, "inferred");
        assert.strictEqual(decision.requiresConfirmation, true);
        assert.strictEqual(decision.enableGoogleSearch, false);
    });

    it("infers web_search for live crypto price queries in auto mode", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [
                {
                    role: "user",
                    content: "what is the bitcoin price and TAO and solana?",
                },
            ],
            docIndex: {},
        });
        assert.strictEqual(decision.turnMode, "web_search");
        assert.strictEqual(decision.enableGoogleSearch, true);
        assert.strictEqual(decision.reason, "web_search_intent");
    });

    it("infers web_search for private stock valuation queries in auto mode", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [
                {
                    role: "user",
                    content: "how much spacex stock trading on?",
                },
            ],
            docIndex: {},
        });
        assert.strictEqual(decision.turnMode, "web_search");
        assert.strictEqual(decision.enableGoogleSearch, true);
        assert.strictEqual(decision.reason, "web_search_intent");
    });

    it("keeps attached-document price questions in normal answer mode", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [
                {
                    role: "user",
                    content: "what is the price of the steel building",
                },
            ],
            docIndex: {
                "doc-0": {
                    document_id: "doc-0",
                    filename: "steel_modified.docx",
                    file_type: "docx",
                },
            },
        });
        assert.strictEqual(decision.turnMode, "answer");
        assert.strictEqual(decision.enableGoogleSearch, false);
        assert.strictEqual(decision.reason, "normal_answer");
    });

    it("downgrades explicit web_search when documents are attached", () => {
        for (const doc of [
            {
                document_id: "doc-0",
                filename: "steel_modified.docx",
                file_type: "docx",
            },
            {
                document_id: "doc-1",
                filename: "contract.pdf",
                file_type: "pdf",
            },
        ]) {
            const decision = resolveTurnMode({
                requestedMode: "web_search",
                messages: [
                    {
                        role: "user",
                        content: "what is the price of the steel building",
                    },
                ],
                docIndex: { "doc-0": doc },
            });
            assert.strictEqual(decision.turnMode, "answer");
            assert.strictEqual(decision.enableGoogleSearch, false);
            assert.strictEqual(decision.reason, "web_search_blocked_docs_attached");
        }
    });

    it("does not infer web_search for explicit web phrasing with attached doc", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [
                {
                    role: "user",
                    content: "search the web for bitcoin",
                },
            ],
            docIndex: {
                "doc-0": {
                    document_id: "doc-0",
                    filename: "steel_modified.docx",
                    file_type: "docx",
                },
            },
        });
        assert.strictEqual(decision.turnMode, "answer");
        assert.strictEqual(decision.enableGoogleSearch, false);
        assert.strictEqual(decision.reason, "normal_answer");
    });

    it("resolveTurnModeAsync keeps steel-building doc questions in answer mode", async () => {
        process.env.SAUDI_LAW_ROUTING = "0";
        process.env.WEB_SEARCH_LLM_ROUTING = "1";
        const decision = await resolveTurnModeAsync({
            requestedMode: "auto",
            messages: [
                {
                    role: "user",
                    content: "what is the price of the steel building",
                },
            ],
            docIndex: {
                "doc-0": {
                    document_id: "doc-0",
                    filename: "steel_modified.docx",
                    file_type: "docx",
                },
            },
        });
        assert.strictEqual(decision.turnMode, "answer");
        assert.strictEqual(decision.enableGoogleSearch, false);
        assert.strictEqual(decision.reason, "normal_answer");
    });

    it("does not infer web_search for ambiguous acronym queries without routing", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [{ role: "user", content: "what is contract law?" }],
            docIndex: {},
        });
        assert.strictEqual(decision.turnMode, "answer");
        assert.strictEqual(decision.enableGoogleSearch, false);
    });

    it("infers web_search for geo-scoped acronym queries via heuristic", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [{ role: "user", content: "what is CJP in india ?" }],
            docIndex: {},
        });
        assert.strictEqual(decision.turnMode, "web_search");
        assert.strictEqual(decision.enableGoogleSearch, true);
        assert.strictEqual(decision.reason, "web_search_intent");
    });

    it("inherits web_search for a short follow-up to a prior live-info turn", async () => {
        process.env.SAUDI_LAW_ROUTING = "0";
        const decision = await resolveTurnModeAsync({
            requestedMode: "auto",
            messages: [
                { role: "user", content: "What is the weather today?" },
                { role: "assistant", content: "It varies by location." },
                { role: "user", content: "for London" },
            ],
            docIndex: {},
            priorTurn: {
                turnMode: "web_search",
                enableGoogleSearch: true,
                reason: "web_search_intent",
            },
        });
        assert.strictEqual(decision.turnMode, "web_search");
        assert.strictEqual(decision.enableGoogleSearch, true);
        assert.strictEqual(decision.reason, "web_search_routing");
        assert.strictEqual(decision.routingContext?.followUpDetected, true);
        assert.strictEqual(decision.routingContext?.inheritedWebSearch, true);
        assert.match(decision.routingContext?.routingText ?? "", /Follow-up: for London/);
    });

    it("does not inherit web_search for a clear new document-generation topic", async () => {
        process.env.SAUDI_LAW_ROUTING = "0";
        const decision = await resolveTurnModeAsync({
            requestedMode: "auto",
            messages: [
                { role: "user", content: "What is the weather today?" },
                { role: "assistant", content: "It varies by location." },
                { role: "user", content: "draft an NDA document" },
            ],
            docIndex: {},
            priorTurn: {
                turnMode: "web_search",
                enableGoogleSearch: true,
                reason: "web_search_intent",
            },
        });
        assert.strictEqual(decision.turnMode, "document_generation");
        assert.strictEqual(decision.enableGoogleSearch, false);
        assert.strictEqual(decision.routingContext?.followUpDetected, false);
    });

    it("maps requestedTrackChanges=false to default_comments_only with auto", () => {
        const decision = resolveTurnMode({
            requestedMode: "auto",
            messages: [{ role: "user", content: "hello" }],
            docIndex: {
                "doc-0": { document_id: "doc-0", filename: "foo.docx", file_type: "docx" }
            },
        });
        assert.strictEqual(decision.turnMode, "answer");
        assert.strictEqual(decision.requiresConfirmation, false);
    });

    it("explicit non-mutating modes override legacy track_changes", () => {
        for (const requestedMode of ["answer", "web_search"] as const) {
            const decision = resolveTurnMode({
                requestedMode,
                requestedTrackChanges: true,
                messages: [{ role: "user", content: "rewrite this document" }],
                docIndex: {
                    "doc-0": {
                        document_id: "doc-0",
                        filename: "foo.docx",
                        file_type: "docx",
                    },
                },
            });
            assert.strictEqual(decision.requestedMode, requestedMode);
            assert.strictEqual(decision.documentMode.trackChanges, false);
            assert.strictEqual(
                isToolAllowedForMode("plan_comments_on_document", decision.turnMode),
                false,
            );
            assert.strictEqual(
                isToolAllowedForMode("edit_document", decision.turnMode),
                false,
            );
        }
    });

    it("enforces the mutation capability matrix independently of prompts", () => {
        assert.strictEqual(isToolAllowedForMode("edit_document", "answer"), false);
        assert.strictEqual(
            isToolAllowedForMode("plan_comments_on_document", "answer"),
            false,
        );
        assert.strictEqual(
            isToolAllowedForMode("plan_comments_on_document", "comments"),
            true,
        );
        assert.strictEqual(isToolAllowedForMode("edit_document", "comments"), false);
        assert.strictEqual(
            isToolAllowedForMode("edit_document", "document_direct_edit"),
            true,
        );
        assert.strictEqual(
            isToolAllowedForMode(
                "plan_comments_on_document",
                "document_direct_edit",
            ),
            false,
        );
        assert.strictEqual(
            isToolAllowedForMode("generate_docx", "document_generation"),
            true,
        );
        assert.strictEqual(
            isToolAllowedForMode("edit_document", "document_generation"),
            false,
        );
        assert.strictEqual(isToolAllowedForMode("edit_document", "web_search"), false);
    });

    it("applySaudiLawRouting forces web_search in auto mode", () => {
        const base = resolveTurnMode({
            requestedMode: "auto",
            messages: [
                {
                    role: "user",
                    content: "What is good faith under the Saudi Civil Transactions Law?",
                },
            ],
            docIndex: {},
        });
        const saudi = buildSaudiLawRoutingResult({
            isSaudiLawQuestion: true,
            source: "heuristic",
            domain: "civil_transactions",
            confidence: "high",
        });
        const decision = applySaudiLawRouting(base, saudi, "auto");
        assert.strictEqual(decision.turnMode, "web_search");
        assert.strictEqual(decision.enableGoogleSearch, true);
        assert.strictEqual(decision.reason, "saudi_law_requires_web_search");
        assert.strictEqual(decision.saudiLawRouting?.domain, "civil_transactions");
    });

    it("applySaudiLawRouting keeps explicit answer mode but enables web search", () => {
        const base = resolveTurnMode({
            requestedMode: "answer",
            messages: [
                {
                    role: "user",
                    content: "Under Saudi law, who is liable in a general partnership?",
                },
            ],
            docIndex: {},
        });
        const saudi = buildSaudiLawRoutingResult({
            isSaudiLawQuestion: true,
            source: "heuristic",
            domain: "companies",
            confidence: "high",
        });
        const decision = applySaudiLawRouting(base, saudi, "answer");
        assert.strictEqual(decision.turnMode, "answer");
        assert.strictEqual(decision.enableGoogleSearch, true);
        assert.strictEqual(decision.reason, "saudi_law_requires_web_search");
    });

    it("applySaudiLawRouting does not override workflow_review", () => {
        const base = resolveTurnMode({
            requestedMode: "auto",
            messages: [
                {
                    role: "user",
                    content: "run the buyer playbook workflow for Saudi law review",
                },
            ],
            docIndex: {},
        });
        const saudi = buildSaudiLawRoutingResult({
            isSaudiLawQuestion: true,
            source: "heuristic",
            domain: "unknown",
            confidence: "high",
        });
        const decision = applySaudiLawRouting(base, saudi, "auto");
        assert.strictEqual(decision.turnMode, "workflow_review");
        assert.strictEqual(decision.enableGoogleSearch, false);
    });

    it("resolveTurnModeAsync routes Saudi heuristic questions to web search", async () => {
        process.env.SAUDI_LAW_ROUTING = "1";
        process.env.SAUDI_LAW_LLM_ROUTING = "0";
        const decision = await resolveTurnModeAsync({
            requestedMode: "auto",
            messages: [
                {
                    role: "user",
                    content:
                        "What company forms are recognized by the Saudi Companies Law?",
                },
            ],
            docIndex: {},
        });
        assert.strictEqual(decision.enableGoogleSearch, true);
        assert.strictEqual(decision.reason, "saudi_law_requires_web_search");
        assert.strictEqual(decision.saudiLawRouting?.isSaudiLawQuestion, true);
    });
});