import { describe, it, expect, beforeAll } from "vitest";
import { config } from "dotenv";
import path from "path";
import { fileURLToPath } from "url";

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);
config({ path: path.resolve(__dirname, "../../.env.local") });

const OLLAMA_BASE_URL = process.env.OLLAMA_BASE_URL;
const OLLAMA_MODEL = process.env.OLLAMA_MODEL;

if (!OLLAMA_BASE_URL || !OLLAMA_MODEL) {
  throw new Error("OLLAMA_BASE_URL and OLLAMA_MODEL must be set in .env.local");
}

/** Make one live Ollama chat call with JSON format. */
async function ollamaCall(instruction, userData) {
  const messages = [
    { role: "system", content: instruction.trim() },
    {
      role: "user",
      content: `Source: ${JSON.stringify(userData.source || "")}

Disagreement:
${(userData.disagreement || []).map((d, i) => `${i + 1}. ${d}`).join("\n")}

Evidence needed:
${(userData.evidenceNeeded || []).map((e, i) => `${i + 1}. ${e}`).join("\n")}

Clarification target: ${userData.clarificationTarget || "N/A"}

Clarification question: ${userData.clarificationQuestion || "N/A"}

User's answer: ${userData.userAnswer || "N/A"}

requiresUserClarification: ${JSON.stringify(userData.requiresUserClarification ?? "N/A")}`,
    },
  ];

  const res = await fetch(`${OLLAMA_BASE_URL}/api/chat`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ model: OLLAMA_MODEL, messages, format: "json", stream: false }),
  });

  if (!res.ok) throw new Error(`Ollama API error: ${res.status} ${res.statusText}`);

  const data = await res.json();
  const rawContent = data.message?.content ?? "";
  const cleaned = rawContent.replace(/```(?:json)?\s*/g, "").replace(/```\s*/g, "");
  return JSON.parse(cleaned.trim());
}

// ──────────────────────────────────────────────
// Evaluation helpers
// ──────────────────────────────────────────────

function evaluateA1(result) {
  const correct = result.requiresUserClarification === true;
  return {
    stage_correct: correct,
    stage_failed: !correct,
    classification: correct ? "stage_correct" : "stage_failed",
    details: `Expected true, got ${result.requiresUserClarification}. Reason: ${result.reason}`,
  };
}

function evaluateA2(result) {
  if (result.clarificationTarget == null) {
    return { stage_failed: true, classification: "stage_failed", details: "Returned null for a user-owned ambiguity." };
  }
  const t = result.clarificationTarget.toLowerCase();
  if (result.clarificationTarget.trim().endsWith("?")) {
    return { stage_failed: true, classification: "stage_failed", details: "Clarification target is worded as a question." };
  }
  const hasRiskOrConstraint = /risk|constraint|preference|trade.?off|boundary/.test(t);
  return {
    stage_correct: hasRiskOrConstraint,
    stage_failed: !hasRiskOrConstraint,
    classification: hasRiskOrConstraint ? "stage_correct" : "stage_failed",
    details: `Target: "${result.clarificationTarget}". ${hasRiskOrConstraint ? "Aligned with risk/preference distinction." : "Does not reference risk/constraint distinction."}`,
  };
}

function evaluateA3(result) {
  if (typeof result.question !== "string" || !result.question.trim()) {
    return { stage_failed: true, classification: "stage_failed", details: "Missing or empty question." };
  }
  const q = result.question;
  if (!q.endsWith("?")) return { stage_failed: true, classification: "stage_failed", details: "Does not end with '?'.", raw: q };
  const qmCount = (q.match(/\?/g) || []).length;
  if (qmCount !== 1) return { stage_failed: true, classification: "stage_failed", details: `Multiple question marks (${qmCount}).`, raw: q };
  const qLower = q.toLowerCase();
  const biasedPhrases = ["you prefer", "your preference is", "rather than the other way around"];
  let hasBias = biasedPhrases.some(bp => qLower.includes(bp));
  const referencesTargetConcept = /risk|constraint|preference|trade.?off/.test(qLower);
  const correct = !hasBias && referencesTargetConcept;
  return {
    stage_correct: correct,
    stage_failed: !correct,
    classification: correct ? "stage_correct" : "stage_failed",
    details: `Question: "${q}". Neutral=${!hasBias}, target-aligned=${referencesTargetConcept}`,
  };
}

function evaluateA4(result) {
  if (typeof result.resolvedMeaning !== "string" || !result.resolvedMeaning.trim()) {
    return { stage_failed: true, classification: "stage_failed", details: "Missing resolvedMeaning." };
  }
  if (typeof result.targetResolved !== "boolean") {
    return { stage_failed: true, classification: "stage_failed", details: "targetResolved is not a boolean." };
  }
  const targetResolvedCorrect = result.targetResolved === true;
  const meaningLower = result.resolvedMeaning.toLowerCase();
  const mentionsConstraint = /constraint|hard constraint|no increase/.test(meaningLower);
  const introducesWiderConsequence = /growth.*impossible|growth.*stopped|which growth option|current risk level/.test(meaningLower);
  const correct = targetResolvedCorrect && mentionsConstraint && !introducesWiderConsequence;
  return {
    stage_correct: correct,
    stage_failed: !correct,
    classification: correct ? "stage_correct" : "stage_failed",
    details: `resolvedMeaning="${result.resolvedMeaning}", targetResolved=${result.targetResolved}, remainingUncertainty=${JSON.stringify(result.remainingUncertainty)}`,
  };
}

// ──────────────────────────────────────────────
// Shared outputs populated in beforeAll
// ──────────────────────────────────────────────

const scenario = { a1: null, a2: null, a3: null, a4: null, timings: {}, b1: null };
let callCount = 0;
const allTimings = [];

// ──────────────────────────────────────────────
// Test suite
// ──────────────────────────────────────────────

describe("Experiment 54W - End-to-End Clarification Chain", () => {
  beforeAll(async () => {
    if (!OLLAMA_BASE_URL || !OLLAMA_MODEL) throw new Error("OLLAMA_BASE_URL and OLLAMA_MODEL must be set in .env.local");
  });

  describe("Scenario A — Full Clarification Journey (4 live calls)", () => {
    it("A1: Resolution Source", async () => {
      const start = Date.now(); callCount++;

      const source = "I want the business to grow, but I don't want to take on more risk.";
      const disagreement = [
        "growth should be prioritised even if some additional risk is unavoidable",
        "avoiding additional risk is a hard constraint even if growth is slower",
      ];
      const evidenceNeeded = ["possible growth opportunities", "risk exposure of each option"];

      const instruction = `Decide whether resolving the stated disagreement requires additional meaning, preference, intent, or factual information that only the user can provide.

Return true when evidence alone cannot settle the disagreement because the missing distinction belongs to the user's intended meaning, priority, constraint, or private knowledge. Return false when the disagreement can be investigated using external, operational, or observable evidence without asking the user to define what they mean.

Return valid JSON only in this shape:
{
  "requiresUserClarification": true | false,
  "reason": "one short sentence"
}

Do not generate a question. Do not choose which interpretation is correct.`;

      scenario.a1 = await ollamaCall(instruction, { source, disagreement, evidenceNeeded });
      scenario.timings.a1 = Date.now() - start;
      allTimings.push({ scenario: "A", stage: "a1", ms: scenario.timings.a1 });

      console.log("\n--- Scenario A Stage A1 (Resolution Source) ---");
      console.log("Actual output:", JSON.stringify(scenario.a1, null, 2));

      expect(scenario.a1.requiresUserClarification).toBe(true);
    }, 120000);

    it("A2: Clarification Target feeds actual A1 output", async () => {
      const start = Date.now(); callCount++;

      const source = "I want the business to grow, but I don't want to take on more risk.";
      const disagreement = [
        "growth should be prioritised even if some additional risk is unavoidable",
        "avoiding additional risk is a hard constraint even if growth is slower",
      ];

      const instruction = `Identify the specific unresolved distinction that only the user can clarify.

If clarification is required (requiresUserClarification: true), return the smallest statement of the missing user-owned meaning, preference, priority, constraint, definition, or private fact.

If clarification is not required (requiresUserClarification: false), return null.

Do not write a question. Do not add evidence needs. Do not select a preferred interpretation.

Return valid JSON only in this shape:
{
  "clarificationTarget": "short statement" | null
}`;

      scenario.a2 = await ollamaCall(instruction, { source, disagreement, requiresUserClarification: scenario.a1.requiresUserClarification });
      scenario.timings.a2 = Date.now() - start;
      allTimings.push({ scenario: "A", stage: "a2", ms: scenario.timings.a2 });

      console.log("\n--- Scenario A Stage A2 (Clarification Target) ---");
      console.log("Input from actual A1:", JSON.stringify(scenario.a1));
      console.log("Actual output:", JSON.stringify(scenario.a2, null, 2));

      expect(scenario.a2.clarificationTarget).toBeDefined();
      expect(typeof scenario.a2.clarificationTarget).toBe("string");
      expect(scenario.a2.clarificationTarget.trim()).not.toMatch(/\?$/);
    }, 120000);

    it("A3: Clarification Question feeds actual A2 output", async () => {
      const start = Date.now(); callCount++;

      const source = "I want the business to grow, but I don't want to take on more risk.";

      const instruction = `Write one concise clarification question that asks only about the supplied clarification target. Keep it neutral between the possible meanings. Do not introduce new facts, assumptions, evidence requests, recommendations, or additional questions. Do not explain why you are asking.

Return valid JSON only in this shape:
{
  "question": "one clarification question"
}`;

      scenario.a3 = await ollamaCall(instruction, { source, clarificationTarget: scenario.a2.clarificationTarget });
      scenario.timings.a3 = Date.now() - start;
      allTimings.push({ scenario: "A", stage: "a3", ms: scenario.timings.a3 });

      console.log("\n--- Scenario A Stage A3 (Clarification Question) ---");
      console.log("Input target from actual A2:", JSON.stringify(scenario.a2.clarificationTarget));
      console.log("Actual output:", JSON.stringify(scenario.a3, null, 2));

      expect(scenario.a3.question).toBeDefined();
      expect(scenario.a3.question.endsWith("?")).toBe(true);
    }, 120000);

    it("A4: Answer Resolution feeds actual A2 target and A3 question", async () => {
      const start = Date.now(); callCount++;

      const source = "I want the business to grow, but I don't want to take on more risk.";
      const userAnswer = "It's a hard constraint. I don't want any increase in risk.";

      const instruction = `Use the user's clarification answer only to resolve the supplied clarification target. State the meaning now established by that answer. Mark targetResolved true only when the answer settles the target. Put any uncertainty that remains specifically about that target into remainingUncertainty; otherwise return null. Do not infer wider consequences, rewrite unrelated source meaning, recommend action, or generate another question.

Return valid JSON only in this shape:
{
  "resolvedMeaning": "short statement",
  "targetResolved": true,
  "remainingUncertainty": null
}`;

      scenario.a4 = await ollamaCall(instruction, { source, clarificationTarget: scenario.a2.clarificationTarget, clarificationQuestion: scenario.a3.question, userAnswer });
      scenario.timings.a4 = Date.now() - start;
      allTimings.push({ scenario: "A", stage: "a4", ms: scenario.timings.a4 });

      console.log("\n--- Scenario A Stage A4 (Answer Resolution) ---");
      console.log("Input target from actual A2:", JSON.stringify(scenario.a2.clarificationTarget));
      console.log("Input question from actual A3:", JSON.stringify(scenario.a3.question));
      console.log("Actual output:", JSON.stringify(scenario.a4, null, 2));

      expect(scenario.a4.resolvedMeaning).toBeDefined();
      expect(typeof scenario.a4.resolvedMeaning).toBe("string");
      expect(scenario.a4.targetResolved).toBe(true);
    }, 120000);
  });

  describe("Scenario B — Evidence-Resolvable Stop (1 live call)", () => {
    it("B1: Resolution Source must return false", async () => {
      const start = Date.now(); callCount++;

      const source = "Orders are arriving late and customers have started complaining.";
      const disagreement = [
        "delays may be caused by insufficient staff capacity",
        "delays may be caused by unreliable supplier lead times",
      ];
      const evidenceNeeded = [
        "staffing levels and workload",
        "processing throughput",
        "supplier lead-time history",
        "supplier delivery reliability",
      ];

      const instruction = `Decide whether resolving the stated disagreement requires additional meaning, preference, intent, or factual information that only the user can provide.

Return true when evidence alone cannot settle the disagreement because the missing distinction belongs to the user's intended meaning, priority, constraint, or private knowledge. Return false when the disagreement can be investigated using external, operational, or observable evidence without asking the user to define what they mean.

Return valid JSON only in this shape:
{
  "requiresUserClarification": true | false,
  "reason": "one short sentence"
}

Do not generate a question. Do not choose which interpretation is correct.`;

      scenario.b1 = await ollamaCall(instruction, { source, disagreement, evidenceNeeded });
      scenario.timings.b1 = Date.now() - start;
      allTimings.push({ scenario: "B", stage: "b1", ms: scenario.timings.b1 });

      console.log("\n--- Scenario B Stage B1 (Resolution Source) ---");
      console.log("Actual output:", JSON.stringify(scenario.b1, null, 2));

      expect(scenario.b1.requiresUserClarification).toBe(false);
      console.log("\n--- Scenario B: chain correctly stopped after B1 (no clarification stages called) ---");
    }, 120000);
  });

  it("54W — aggregate results", () => {
    const a1Eval = evaluateA1(scenario.a1);
    const a2Eval = evaluateA2(scenario.a2);
    const a3Eval = evaluateA3(scenario.a3);
    const a4Eval = evaluateA4(scenario.a4);
    const b1Correct = scenario.b1.requiresUserClarification === false;

    const allStagesCorrect = a1Eval.stage_correct && a2Eval.stage_correct && a3Eval.stage_correct && a4Eval.stage_correct;
    const chainResult = allStagesCorrect ? "chain_correct" : "chain_drifted";

    // Determine first drift point
    let firstDrift = null;
    if (!a1Eval.stage_correct) firstDrift = "A1 (Resolution Source)";
    else if (!a2Eval.stage_correct) firstDrift = "A2 (Clarification Target)";
    else if (!a3Eval.stage_correct) firstDrift = "A3 (Clarification Question)";
    else if (!a4Eval.stage_correct) firstDrift = "A4 (Answer Resolution)";

    const totalTime = allTimings.reduce((s, t) => s + t.ms, 0);

    console.log("\n=== Experiment 54W Summary ===");
    console.log(`Total live calls: ${allTimings.length}`);
    console.log(`Total inference time: ${totalTime}ms`);
    console.log(`Average time: ${(totalTime / allTimings.length).toFixed(1)}ms per call`);
    console.log(`Fastest: ${Math.min(...allTimings.map((t) => t.ms))}ms`);
    console.log(`Slowest: ${Math.max(...allTimings.map((t) => t.ms))}ms`);

    console.log("\n--- Scenario A Stage-by-Stage ---");
    console.log(`A1: ${a1Eval.classification} — ${a1Eval.details}`);
    console.log(`A2: ${a2Eval.classification} — ${a2Eval.details}`);
    console.log(`A3: ${a3Eval.classification} — ${a3Eval.details}`);
    console.log(`A4: ${a4Eval.classification} — ${a4Eval.details}`);

    // Print actual upstream outputs that were passed downstream
    console.log("\n--- Actual Upstream Outputs Used Between Stages ---");
    console.log(`A2 received from A1: requiresUserClarification=${scenario.a1.requiresUserClarification}, reason="${scenario.a1.reason}"`);
    console.log(`A3 received from A2: clarificationTarget="${scenario.a2.clarificationTarget}"`);
    console.log(`A4 received from A2: clarificationTarget="${scenario.a2.clarificationTarget}"`);
    console.log(`A4 received from A3: question="${scenario.a3.question}"`);

    // Check for drift accumulation and unsupported meaning
    let introducesUnsupported = false;
    if (scenario.a4) {
      const meaningLower = scenario.a4.resolvedMeaning?.toLowerCase() || "";
      introducesUnsupported = /growth.*impossible|growth.*stopped|which growth option|current risk level/.test(meaningLower);
    }

    console.log("\n--- Scenario B ---");
    console.log(`B1: ${b1Correct ? "correct_stop" : "incorrect_clarification"} — ${scenario.b1.details || scenario.b1.reason}`);

    console.log("\n=== Overall Results ===");
    console.log(`Scenario A chain: ${chainResult}`);
    if (firstDrift) console.log(`First drift point: ${firstDrift}`);
    else console.log("No drift detected.");
    console.log(`Scenario B: ${b1Correct ? "correct_stop" : "incorrect_clarification"}`);
    console.log(`Did any stage choose a winner: ${!allStagesCorrect || "No (checked via output invariants)"}`);

    expect(allTimings.length).toBeLessThanOrEqual(5);
  });
});