#!/usr/bin/env python3
from __future__ import annotations

import json
import os
import sys
import unittest
from pathlib import Path
from unittest import mock


HERE = Path(__file__).resolve().parent

os.environ["ANTHROPIC_API_KEY"] = "test-key"

if str(HERE) not in sys.path:
    sys.path.insert(0, str(HERE))

import server as mcp_server  # noqa: E402


VALID_PLAN = {
    "meta": {
        "title": "Founder Fortune Cookie",
        "description": "A playful interactive prompt.",
        "document_format": "Narrative Scroll",
    },
    "theme": "obsidian-cyan",
    "wrapper": {"engine": "01"},
    "sections": [
        {
            "title": "Cookie",
            "content": {
                "headline": "Crack the cookie",
                "body": "A tiny interactive joke for founders.",
            },
        }
    ],
}


class FreeformPlannerProgressTests(unittest.IsolatedAsyncioTestCase):
    async def test_call_planner_reports_quality_retry_progress(self) -> None:
        progress_events: list[tuple[int, str]] = []
        audit_calls = {"count": 0}

        async def fake_call_anthropic(system: str, user: str, **kwargs):
            step_name = kwargs.get("step_name")
            if step_name == "planner":
                payload = {
                    "content": [{"type": "text", "text": json.dumps(VALID_PLAN)}],
                    "usage": {"input_tokens": 10, "output_tokens": 20},
                }
            elif step_name == "planner_quality_retry":
                payload = {
                    "content": [{"type": "text", "text": json.dumps(VALID_PLAN)}],
                    "usage": {"input_tokens": 8, "output_tokens": 18},
                }
            else:
                raise AssertionError(f"Unexpected planner step: {step_name}")
            return payload, {"step_name": step_name}

        def fake_audit_plan_against_prompt_stack(**kwargs):
            audit_calls["count"] += 1
            if audit_calls["count"] == 1:
                return {
                    "score_penalty": 6,
                    "critical": False,
                    "should_retry": True,
                    "feedback_lines": ["Use a more playful, founder-friendly structure."],
                    "warnings": [],
                }
            return {
                "score_penalty": 0,
                "critical": False,
                "should_retry": False,
                "feedback_lines": [],
                "warnings": [],
            }

        with mock.patch.object(mcp_server.SETTINGS, "anthropic_api_key", "test-key"), mock.patch.object(
            mcp_server, "call_anthropic", side_effect=fake_call_anthropic
        ), mock.patch.object(
            mcp_server, "audit_plan_against_prompt_stack", side_effect=fake_audit_plan_against_prompt_stack
        ), mock.patch.object(
            mcp_server, "resolve_active_prompt_modules", return_value=[]
        ):
            plan, payload = await mcp_server.call_planner(
                prompt="Build a fun founder fortune cookie artifact.",
                progress_callback=lambda progress, message: progress_events.append((progress, message)),
            )

        self.assertEqual(plan["meta"]["title"], "Founder Fortune Cookie")
        self.assertEqual([event[0] for event in progress_events], [20, 26])
        self.assertIn("corrective retry", progress_events[1][1].lower())
        self.assertEqual(payload.get("retry_strategy"), "quality-critic")

    async def test_call_planner_reports_original_retry_and_json_repair_progress(self) -> None:
        progress_events: list[tuple[int, str]] = []

        async def fake_call_anthropic(system: str, user: str, **kwargs):
            step_name = kwargs.get("step_name")
            if step_name in {"planner", "planner_retry"}:
                payload = {
                    "content": [{"type": "text", "text": '{"meta": {"title": "Broken Plan"'}],
                    "usage": {"input_tokens": 4, "output_tokens": 4},
                }
            elif step_name == "json_repair":
                payload = {
                    "content": [{"type": "text", "text": json.dumps(VALID_PLAN)}],
                    "usage": {"input_tokens": 5, "output_tokens": 7},
                }
            else:
                raise AssertionError(f"Unexpected planner step: {step_name}")
            return payload, {"step_name": step_name}

        with mock.patch.object(mcp_server.SETTINGS, "anthropic_api_key", "test-key"), mock.patch.object(
            mcp_server, "call_anthropic", side_effect=fake_call_anthropic
        ), mock.patch.object(
            mcp_server, "resolve_active_prompt_modules", return_value=[]
        ):
            plan, payload = await mcp_server.call_planner(
                prompt="Build a fun founder fortune cookie artifact.",
                progress_callback=lambda progress, message: progress_events.append((progress, message)),
            )

        self.assertEqual(plan["meta"]["title"], "Founder Fortune Cookie")
        self.assertEqual([event[0] for event in progress_events], [20, 28, 32])
        self.assertIn("invalid json", progress_events[1][1].lower())
        self.assertIn("attempting json repair", progress_events[2][1].lower())
        self.assertEqual(payload.get("retry_strategy"), "json-repair")


if __name__ == "__main__":
    unittest.main()
