from __future__ import annotations

import asyncio
from pathlib import Path
import sys
import types

import pytest


def _install_acp_stubs():
    acp_module = types.ModuleType("acp")
    acp_module.PROTOCOL_VERSION = 1
    acp_module.Client = object
    acp_module.RequestError = Exception
    acp_module.connect_to_agent = object
    acp_module.text_block = lambda text: text
    acp_module.image_block = lambda **kwargs: {"type": "image", **kwargs}
    sys.modules.setdefault("acp", acp_module)

    core_module = types.ModuleType("acp.core")
    core_module.ClientSideConnection = object
    sys.modules.setdefault("acp.core", core_module)

    schema_module = types.ModuleType("acp.schema")
    for name in (
        "AgentMessageChunk",
        "AgentPlanUpdate",
        "AgentThoughtChunk",
        "AllowedOutcome",
        "AvailableCommandsUpdate",
        "ClientCapabilities",
        "CurrentModeUpdate",
        "DeniedOutcome",
        "EmbeddedResourceContentBlock",
        "FileSystemCapabilities",
        "PermissionOption",
        "ReadTextFileResponse",
        "RequestPermissionResponse",
        "ResourceContentBlock",
        "SessionInfoUpdate",
        "TextContentBlock",
        "ToolCall",
        "ToolCallProgress",
        "ToolCallStart",
        "UsageUpdate",
        "UserMessageChunk",
        "WriteTextFileResponse",
    ):
        setattr(schema_module, name, type(name, (), {}))
    sys.modules.setdefault("acp.schema", schema_module)


_install_acp_stubs()

from investibook import config as cfg  # noqa: E402
from investibook.api import chat  # noqa: E402
from investibook.core.agent_runtime import AgentSession, _workspace_id_from_dir  # noqa: E402
from investibook.core import workspace_evidence_postpass as postpass  # noqa: E402
from investibook.db import queries as dbq  # noqa: E402


def _init_tmp_db(tmp_path, monkeypatch):
    db_path = tmp_path / "investibook.db"
    workspaces = tmp_path / "workspaces"
    workspaces.mkdir()
    monkeypatch.setattr(dbq, "DB_PATH", db_path)
    monkeypatch.setattr(cfg, "DB_PATH", db_path)
    monkeypatch.setattr(cfg, "WORKSPACES_DIR", workspaces)
    dbq._db_init()
    return workspaces


def test_acp_workspace_identity_accepts_lark_workspace_ids():
    lark_dir = Path("/tmp/workspaces/lark-ws-5c9bbb67-7f87-4f21-acf8-e646fb0551dc")
    web_dir = Path("/tmp/workspaces/ws-1778486510665-211c5da9f")

    assert _workspace_id_from_dir(lark_dir) == lark_dir.name
    assert _workspace_id_from_dir(web_dir) == web_dir.name
    assert _workspace_id_from_dir(Path("/tmp/workspaces/not-a-workspace")) is None

    session = AgentSession(
        frontend_id="frontend-session",
        agent_type="opencode",
        workspace_dir=lark_dir,
    )
    assert session.workspace_id == lark_dir.name


class _FakeClient:
    def __init__(self) -> None:
        self.queue: asyncio.Queue[dict] = asyncio.Queue()
        self._msg_id: str | None = None
        self.metrics_user_email = None
        self.metrics_channel = None
        self.metrics_usage_scenario = None

    def start_assistant_message(self, message_id: str) -> None:
        self._msg_id = message_id

    def get_and_reset_usage(self) -> dict:
        return {
            "input_tokens": 0,
            "output_tokens": 0,
            "reasoning_tokens": 0,
            "cache_read_tokens": 0,
            "cache_write_tokens": 0,
            "context_used": 0,
            "context_size": 0,
            "cost_amount": 0.0,
            "cost_currency": "USD",
        }


class _FakeSession:
    def __init__(self) -> None:
        self.frontend_id = "frontend-session"
        self.runtime_session_id = "acp-session"
        self.model_tier = "standard"
        self._prompt_lock = asyncio.Lock()
        self._current_prompt_task = None
        self.client = _FakeClient()

    async def prompt(self, message: str, user_email: str | None = None) -> None:
        await self.client.queue.put({"type": "text", "delta": f"echo:{message}"})

    async def persist_assistant_message_snapshot(self) -> None:
        return None


def test_run_chat_and_wait_starts_workspace_evidence_postpass(monkeypatch):
    async def _run():
        session = _FakeSession()
        calls: list[dict] = []

        async def _fake_get_or_create_session(*args, **kwargs):
            return session

        async def _fake_cost_gate(*args, **kwargs):
            return None

        def _fake_start_postpass(**kwargs):
            calls.append(kwargs)

        monkeypatch.setattr(chat, "get_or_create_session", _fake_get_or_create_session)
        monkeypatch.setattr(chat, "effective_prompt_after_cost_gate", _fake_cost_gate)
        monkeypatch.setattr(chat, "_db_get_blob", lambda *args, **kwargs: [])
        monkeypatch.setattr(chat, "_db_get_user_by_email", lambda *args, **kwargs: {"id": "u-1"})
        monkeypatch.setattr(chat, "_start_workspace_evidence_postpass", _fake_start_postpass)

        frontend_id, result = await chat.run_chat_and_wait(
            message="check qps",
            workspace_id="ws-test",
            user_email="dev@localhost",
            timeout_s=5,
            idle_timeout_s=5,
            metrics_user_email="dev@localhost",
            metrics_channel="test",
        )

        assert frontend_id == "frontend-session"
        assert result.finish_reason == "completed"
        assert result.text == "echo:check qps"
        assert len(calls) == 1
        assert calls[0]["workspace_id"] == "ws-test"
        assert calls[0]["user_email"] == "dev@localhost"
        assert calls[0]["assistant_message_id"].startswith("assistant-")

    asyncio.run(_run())


def test_workspace_evidence_postpass_merges_analysis(monkeypatch):
    async def _run():
        llm_calls: list[str] = []
        merge_calls: list[dict] = []

        def _fake_state(workspace_id: str, *, message_id: str):
            assert workspace_id == "ws-test"
            assert message_id == "assistant-1"
            return {
                "needs_postpass": True,
                "report_path": "",
                "report": {
                    "metadata": {"cluster_id": "cluster-1"},
                    "evidence": {"metrics": [{"name": "TiDB Total QPS", "metric_group_id": "tidb_total_qps"}]},
                    "analysis": {},
                },
            }

        async def _fake_llm(prompt: str):
            llm_calls.append(prompt)
            return {
                "add_key_evidence": [
                    {
                        "name": "QPS stable",
                        "type": "metric",
                        "reason": "QPS evidence supports the answer.",
                    }
                ]
            }

        def _fake_update(workspace_id: str, patch: dict, **kwargs):
            merge_calls.append(
                {
                    "workspace_id": workspace_id,
                    "patch": patch,
                    "kwargs": kwargs,
                }
            )
            return {"analysis_fields": sorted(patch.keys())}

        monkeypatch.setattr(postpass, "get_workspace_evidence_postpass_state", _fake_state)
        monkeypatch.setattr(postpass, "_call_direct_llm", _fake_llm)
        monkeypatch.setattr(postpass, "update_workspace_evidence_analysis", _fake_update)

        await postpass.run_workspace_evidence_postpass(
            workspace_id="ws-test",
            message_id="assistant-1",
        )

        assert len(llm_calls) == 1
        assert "Evidence summary for postpass:" in llm_calls[0]

        assert len(merge_calls) == 1
        assert merge_calls[0]["workspace_id"] == "ws-test"
        assert merge_calls[0]["patch"] == {
            "add_key_evidence": [
                {
                    "name": "QPS stable",
                    "type": "metric",
                    "reason": "QPS evidence supports the answer.",
                    "metric_ref": {"metric_group_id": "tidb_total_qps"},
                }
            ]
        }
        assert merge_calls[0]["kwargs"]["cluster_id"] == "cluster-1"
        assert merge_calls[0]["kwargs"]["source"] == "backend-postpass"
        assert merge_calls[0]["kwargs"]["message_id"] == "assistant-1"

    asyncio.run(_run())


def test_workspace_evidence_postpass_records_direct_llm_usage(tmp_path, monkeypatch):
    _init_tmp_db(tmp_path, monkeypatch)

    async def _run():
        def _fake_state(workspace_id: str, *, message_id: str):
            return {
                "needs_postpass": True,
                "report_path": "",
                "report": {
                    "metadata": {"cluster_id": "cluster-1"},
                    "evidence": {"metrics": [{"name": "TiDB Total QPS", "metric_group_id": "tidb_total_qps"}]},
                    "analysis": {},
                },
            }

        async def _fake_llm(prompt: str):
            return (
                {
                    "add_key_evidence": [
                        {
                            "name": "QPS stable",
                            "type": "metric",
                            "reason": "QPS evidence supports the answer.",
                        }
                    ]
                },
                {"prompt_tokens": 11, "completion_tokens": 22},
            )

        def _fake_update(workspace_id: str, patch: dict, **kwargs):
            return {"analysis_fields": sorted(patch.keys())}

        monkeypatch.setattr(postpass, "get_workspace_evidence_postpass_state", _fake_state)
        monkeypatch.setattr(postpass, "_call_direct_llm", _fake_llm)
        monkeypatch.setattr(postpass, "update_workspace_evidence_analysis", _fake_update)

        await postpass.run_workspace_evidence_postpass(
            workspace_id="ws-test",
            message_id="assistant-usage",
        )

    asyncio.run(_run())

    usage = dbq._db_get_workspace_token_usage("ws-test")
    assert len(usage) == 1
    assert usage[0]["model_tier"] == "evidence-postpass"
    assert usage[0]["channel"] == "evidence-postpass"
    assert usage[0]["session_id"] == "evidence-postpass:assistant-usage"
    assert usage[0]["input_tokens"] == 11
    assert usage[0]["output_tokens"] == 22
    assert usage[0]["cost_amount"] == pytest.approx(0.000363)
    assert usage[0]["cost_currency"] == "USD"


def test_workspace_evidence_postpass_adds_missing_metric_ref_by_name():
    patch = postpass._normalize_analysis_patch_refs(
        {
            "add_key_evidence": [
                {
                    "name": "TiDB Total QPS",
                    "type": "metric",
                    "reason": "QPS changed during the incident window.",
                }
            ]
        },
        {
            "evidence": {
                "metrics": [
                    {
                        "name": "TiDB Total QPS",
                        "metric_group_id": "tidb_total_qps",
                        "summary": {"point_count": 120},
                    }
                ]
            }
        },
    )

    assert patch["add_key_evidence"][0]["metric_ref"] == {"metric_group_id": "tidb_total_qps"}


def test_workspace_evidence_postpass_skips_when_llm_returns_empty_patch(monkeypatch):
    async def _run():
        merge_calls: list[dict] = []

        def _fake_state(workspace_id: str, *, message_id: str):
            return {
                "needs_postpass": True,
                "report_path": "",
                "report": {
                    "metadata": {"cluster_id": "cluster-1"},
                    "evidence": {
                        "metrics": [
                            {
                                "name": "TiDB Total QPS",
                                "metric_group_id": "tidb_qps_total",
                                "summary": {"point_count": 121},
                            }
                        ]
                    },
                    "analysis": {},
                },
            }

        async def _fake_llm(prompt: str):
            return {}

        def _fake_update(workspace_id: str, patch: dict, **kwargs):
            merge_calls.append({"workspace_id": workspace_id, "patch": patch, "kwargs": kwargs})
            return {"analysis_fields": sorted(patch.keys())}

        monkeypatch.setattr(postpass, "get_workspace_evidence_postpass_state", _fake_state)
        monkeypatch.setattr(postpass, "_call_direct_llm", _fake_llm)
        monkeypatch.setattr(postpass, "update_workspace_evidence_analysis", _fake_update)

        await postpass.run_workspace_evidence_postpass(
            workspace_id="ws-test",
            message_id="assistant-1",
        )

        assert merge_calls == []

    asyncio.run(_run())


def test_workspace_evidence_postpass_skips_when_llm_patch_is_invalid(monkeypatch):
    async def _run():
        merge_calls: list[dict] = []

        def _fake_state(workspace_id: str, *, message_id: str):
            return {
                "needs_postpass": True,
                "report_path": "",
                "report": {
                    "metadata": {"cluster_id": "cluster-1"},
                    "evidence": {
                        "metrics": [
                            {
                                "name": "TiDB Total QPS",
                                "metric_group_id": "tidb_qps_total",
                                "summary": {"point_count": 121},
                            }
                        ]
                    },
                    "analysis": {},
                },
            }

        async def _fake_llm(prompt: str):
            return {"add_key_evidence": {"type": "metric"}}

        def _fake_update(workspace_id: str, patch: dict, **kwargs):
            if not isinstance(patch.get("add_key_evidence"), list):
                raise SystemExit("analysis.add_key_evidence must be a list")
            merge_calls.append({"workspace_id": workspace_id, "patch": patch, "kwargs": kwargs})
            return {"analysis_fields": sorted(patch.keys())}

        monkeypatch.setattr(postpass, "get_workspace_evidence_postpass_state", _fake_state)
        monkeypatch.setattr(postpass, "_call_direct_llm", _fake_llm)
        monkeypatch.setattr(postpass, "update_workspace_evidence_analysis", _fake_update)

        await postpass.run_workspace_evidence_postpass(
            workspace_id="ws-test",
            message_id="assistant-1",
        )

        assert merge_calls == []

    asyncio.run(_run())


def test_workspace_evidence_postpass_fallback_patch_from_metrics():
    patch = postpass._fallback_patch_from_report(
        {
            "evidence": {
                "metrics": [
                    {
                        "name": "TiDB Total QPS",
                        "metric_group_id": "tidb_qps_total",
                        "summary": {"point_count": 121},
                    },
                    {
                        "name": "TiDB QPS by instance",
                        "metric_group_id": "tidb_qps_by_instance",
                        "summary": {"point_count": 1815},
                    }
                ]
            }
        }
    )
    assert patch == {}


def test_workspace_evidence_postpass_prompt_uses_summary_not_report_payload():
    prompt = postpass._build_workspace_evidence_postpass_prompt(
        {
            "metadata": {"cluster_id": "cluster-1"},
            "links": {"clinic_overview": "https://clinic.example"},
            "evidence": {
                "metrics": [
                    {
                        "name": "Huge raw payload",
                        "metric_group_id": "huge_raw_payload",
                        "chart_data": {"series": [{"values": [[1, "2"]]}]},
                    }
                ]
            },
        }
    )
    assert "Evidence summary for postpass:" in prompt
    assert "metric_group_id=huge_raw_payload" in prompt
    assert "chart_data" not in prompt


def test_workspace_evidence_postpass_summary_includes_direct_cli_metrics():
    summary = postpass._build_postpass_summary(
        {
            "metadata": {"cluster_id": "cluster-1", "time_range": {"start": 1, "end": 2}},
            "evidence": {
                "metrics": [
                    {
                        "name": "TiDB QPS",
                        "metric_group_id": "tidb_qps",
                        "promql": "sum(rate(tidb_server_query_total[1m]))",
                        "summary": {"point_count": 120, "max": 99.5},
                        "chart_data": {"series": [{"values": [[1, "2"]]}]},
                    }
                ]
            },
            "links": {"clinic_overview": "https://clinic.example"},
        }
    )
    assert "metric_group_id=tidb_qps" in summary
    assert "summary_point_count=120" in summary
    assert "summary_max=99.5" in summary
    assert "chart_data" not in summary