import asyncio
import uuid
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

from claude_agent_sdk import (
    AssistantMessage,
    ResultMessage,
    TextBlock,
    ToolResultBlock,
    ToolUseBlock,
    UserMessage,
    project_key_for_directory,
)
from test_conversations import ScriptedClient, World, world

from beaver_gateway.agents.claude import ClaudeAgent, ClaudeOptions
from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend
from beaver_gateway.core.conversations import ConversationTexts
from beaver_gateway.core.distill import (
    Distiller,
    DistillContext,
    LineCap,
    check_digest,
    trim_summary,
)
from beaver_gateway.core.gateway_tools import _tools, build_tool_server
from beaver_gateway.core.registry import AgentRegistry
from beaver_gateway.core.scheduler import Job, JobRun, Scheduler
from beaver_gateway.core.transcript import build_entries
from beaver_gateway.storage.models import Conversation

__all__ = ["world"]

DIGEST = """---
type: выжимка
source: "[[{chat}]]"
date: 2026-08-29
---
# тема

## контекст
зачем открывали

## решили
- одно
"""


class DistillerClient(ScriptedClient):
    """The distiller: writes a digest file when told to, answers in N lines."""

    digest_dir: Path | None = None
    lines = 3
    write = True
    frontmatter = DIGEST

    async def receive_response(self):
        prompt = self.prompts[-1]
        written: list[ToolUseBlock] = []
        if self.write and self.digest_dir is not None and "файл не пиши" not in prompt:
            chat = prompt.split("«", 1)[1].split("»", 1)[0] if "«" in prompt else "чат"
            path = self.digest_dir / "2026-08-29 - тема.md"
            path.write_text(self.frontmatter.format(chat=chat), encoding="utf-8")
            written.append(
                ToolUseBlock(id="w1", name="Write", input={"file_path": str(path)})
            )
        for block in written:
            yield AssistantMessage(content=[block], model="m")
            yield UserMessage(content=[ToolResultBlock(tool_use_id=block.id)])
        text = "\n".join(f"строка {i} про {prompt[:20]!r}" for i in range(self.lines))
        yield AssistantMessage(content=[TextBlock(text=text)], model="m")
        yield ResultMessage(
            subtype="success",
            duration_ms=1,
            duration_api_ms=1,
            is_error=False,
            num_turns=1,
            session_id=self.session_id,
            stop_reason="end_turn",
            total_cost_usd=0.0,
            usage={"input_tokens": 1, "output_tokens": 1},
        )


class ClosingClient(ScriptedClient):
    """A deep chat that calls ``close_chat`` inside its reply."""

    conversations: Any = None
    key = ""

    async def receive_response(self):
        prompt = self.prompts[-1]
        if "обсудили" in prompt:
            assert "gateway" in self.options.mcp_servers
            tools = _tools(ClosingClient.conversations, ClosingClient.key)
            close = next(t for t in tools if t.name == "close_chat")
            result = await close.handler({})
            assert "closes after this reply" in result["content"][0]["text"]
            yield AssistantMessage(
                content=[
                    ToolUseBlock(id="c1", name="mcp__gateway__close_chat", input={})
                ],
                model="m",
            )
        yield AssistantMessage(content=[TextBlock(text=f"ok:{prompt}")], model="m")
        yield ResultMessage(
            subtype="success",
            duration_ms=1,
            duration_api_ms=1,
            is_error=False,
            num_turns=1,
            session_id=self.session_id,
            stop_reason="end_turn",
            total_cost_usd=0.0,
            usage={"input_tokens": 1, "output_tokens": 1},
        )


class CapClient(ScriptedClient):
    """A job that rewrites the capped file with ``lines`` lines."""

    target: Path | None = None
    lines = 70

    async def receive_response(self):
        assert self.target is not None
        self.target.write_text(
            "\n".join(f"- строка {i}" for i in range(self.lines)), encoding="utf-8"
        )
        CapClient.lines = 10
        async for event in super().receive_response():
            yield event


def distiller(world: World, client: type[ScriptedClient]) -> Distiller:
    vault = world.root / "vault"
    digests = vault / "мета" / "бобер" / "выжимки"
    digests.mkdir(parents=True)
    agent = ClaudeAgent(
        name="x",
        model="m",
        system_prompt="distill",
        cwd=vault,
        kinds=("fork", "job"),
        options=ClaudeOptions(
            effort="medium", tools=("Read", "Write"), include_partial_messages=False
        ),
    )
    backend = ClaudeSdkBackend(
        agent=agent,
        mcp_internal_urls={},
        session_store=world.store,
        client_factory=client,
        work_dir=world.root / "work",
        pool=world.pool,
    )
    world.conversations._agents = AgentRegistry(  # noqa: SLF001
        [world.agent, world.deep_agent, agent]
    )
    world.conversations._backends["x"] = backend  # noqa: SLF001
    config = Distiller(agent="x", dir=digests, index=digests.parent / "индекс.md")
    world.conversations._distiller = config  # noqa: SLF001
    DistillerClient.digest_dir = digests
    DistillerClient.lines = 3
    DistillerClient.write = True
    DistillerClient.frontmatter = DIGEST
    return config


async def deep_chat(world: World, name: str = "2026-08-20 - тема чата") -> Conversation:
    sid = str(uuid.uuid4())
    await world.store.append(
        world.key(sid),
        build_entries(
            [
                {"role": "user", "content": "q"},
                {
                    "role": "assistant",
                    "content": [
                        {"type": "tool_use", "id": "t1", "name": "Read", "input": {}}
                    ],
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "tool_result", "tool_use_id": "t1", "content": "x"}
                    ],
                },
                {"role": "assistant", "content": "a"},
            ],
            session_id=sid,
            cwd=str(world.root),
            model="m",
        ),
    )
    conv = await world.conversations.create(
        kind="deep", agent="d", origin="markdown", session_id=sid
    )
    await world.conversations.bind(
        conv, frontend="markdown", external_id=f"2026-08/{name}.md"
    )
    return conv


async def test_distill_writes_the_digest_indexes_it_and_merges_short(
    world: World,
) -> None:
    config = distiller(world, DistillerClient)
    prompts: list[DistillContext] = []

    def distill_prompt(ctx: DistillContext) -> str:
        prompts.append(ctx)
        return f"Чат «{ctx.chat_name}» закрыт ({ctx.reason})."

    world.conversations._texts = ConversationTexts(distill=distill_prompt)  # noqa: SLF001
    master = await world.conversations.create(kind="master", agent="a", origin="test")
    chat = await deep_chat(world)

    result = await world.conversations.distill(chat, reason="api")

    assert prompts[0].chat_name == "2026-08-20 - тема чата"
    assert prompts[0].memory is True
    assert result.digest is not None and result.error is None
    assert result.digest.path == config.dir / "2026-08-29 - тема.md"
    assert result.digest.source == "[[2026-08-20 - тема чата]]"
    index = config.index.read_text(encoding="utf-8")
    assert index.startswith("# индекс")
    assert "- 2026-08-29 [[2026-08-20 - тема чата]] → [[2026-08-29 - тема]]" in index
    assert result.text.count("\n") == 2 and not result.trimmed
    closed = await world.conversations.get(chat.external_id)
    assert closed.status == "closed"
    assert closed.flags["digest"] == str(result.digest.path)
    assert closed.flags["closed_reason"] == "api"
    fork = await world.conversations.get(result.fork.external_id)
    assert fork.kind == "fork" and fork.agent_name == "x" and fork.status == "closed"
    items = await world.conversations.queue.recent(master.id)
    assert items[0].origin == "выжимка"
    assert items[0].text.startswith(
        "Закрыт глубокий чат [[2026-08-20 - тема чата]], выжимка [[2026-08-29 - тема]]."
    )
    assert items[0].text.endswith(result.text)
    forked = ScriptedClient.instances[-1]
    assert forked.options.resume is not None
    vault_key = {
        "project_key": project_key_for_directory(str(config.dir.parents[2])),
        "session_id": forked.options.resume,
    }
    entries = await world.store.load(vault_key)
    assert entries
    assert all(
        not isinstance(e.get("message", {}).get("content"), list)
        or all(b.get("type") == "text" for b in e["message"]["content"])
        for e in entries
    )


async def test_memory_off_merges_without_a_file(world: World) -> None:
    config = distiller(world, DistillerClient)
    master = await world.conversations.create(kind="master", agent="a", origin="test")
    chat = await deep_chat(world)
    await world.conversations.set_flags(chat, {"memory": False})

    result = await world.conversations.distill(chat, reason="api")

    assert result.digest is None and result.error is None
    assert list(config.dir.iterdir()) == []
    assert not config.index.exists()
    assert (await world.conversations.get(chat.external_id)).status == "closed"
    items = await world.conversations.queue.recent(master.id)
    assert len(items) == 1 and items[0].text.endswith(result.text)
    assert ", выжимка" not in items[0].text
    assert "файл не пиши" in ScriptedClient.instances[-1].prompts[0]


async def test_bad_frontmatter_and_long_merge_are_reported(world: World) -> None:
    distiller(world, DistillerClient)
    DistillerClient.frontmatter = "---\ntype: заметка\n---\n# тема\n"
    DistillerClient.lines = 9
    await world.conversations.create(kind="master", agent="a", origin="test")
    chat = await deep_chat(world)

    result = await world.conversations.distill(chat, reason="idle 2d")

    assert result.digest is None
    assert result.error is not None and "`type`" in result.error
    assert result.trimmed and result.text.count("\n") == 4
    closed = await world.conversations.get(chat.external_id)
    assert closed.status == "closed"
    assert closed.flags["digest_error"] == result.error


def test_check_digest_rejects_what_is_not_a_digest(tmp_path: Path) -> None:
    config = Distiller(agent="x", dir=tmp_path, index=tmp_path / "i.md")
    path = tmp_path / "d.md"
    path.write_text("---\ntype: выжимка\nsource: ''\ndate: 2026-08-29\n---\nx\n")
    assert check_digest(path, config) == "`source` пустой"
    path.write_text("---\ntype: выжимка\nsource: '[[a]]'\ndate: вчера\n---\nx\n")
    assert "`date`" in check_digest(path, config)
    path.write_text("---\ntype: выжимка\nsource: '[[a]]'\ndate: 2026-08-29\n---\n\n")
    assert check_digest(path, config) == "тело пустое"
    path.write_text("---\ntype: выжимка\nsource: '[[a]]'\ndate: 2026-08-29\n---\nx\n")
    digest = check_digest(path, config)
    assert not isinstance(digest, str) and digest.date.isoformat() == "2026-08-29"
    assert trim_summary("a\n\nb\nc\nd\ne\nf") == ("a\nb\nc\nd\ne", True)


async def test_close_chat_tool_closes_after_the_reply(world: World) -> None:
    distiller(world, DistillerClient)
    world.deep_agent = ClaudeAgent(
        name="d",
        model="m",
        system_prompt="deep",
        cwd=world.root,
        gateway_tools=("close_chat",),
    )
    world.conversations._agents = AgentRegistry(  # noqa: SLF001
        [world.agent, world.deep_agent, world.conversations._agents.get("x")]  # noqa: SLF001
    )
    world.conversations._backends["d"] = ClaudeSdkBackend(  # noqa: SLF001
        agent=world.deep_agent,
        mcp_internal_urls={},
        session_store=world.store,
        client_factory=ClosingClient,
        work_dir=world.root / "work",
        pool=world.pool,
        tool_server=lambda key, _kind: build_tool_server(
            world.conversations, conversation_key=key, names=("close_chat",)
        ),
    )
    master = await world.conversations.create(kind="master", agent="a", origin="test")
    chat = await deep_chat(world)
    ClosingClient.conversations = world.conversations
    ClosingClient.key = chat.external_id

    await world.conversations.post(chat, "ок, обсудили")
    await world.settle(chat, 1)
    for _ in range(100):
        row = await world.conversations.get(chat.external_id)
        if row.status == "closed":
            break
        await asyncio.sleep(0.05)
    else:
        raise AssertionError("chat did not close")

    assert row.flags["closed_reason"] == "close_chat"
    assert row.flags["digest"] is not None
    items = await world.conversations.queue.recent(master.id)
    assert items[0].origin == "выжимка"


async def test_idle_picks_quiet_chats_after_launch_at_most_limit(world: World) -> None:
    distiller(world, DistillerClient)
    now = datetime.now(UTC)
    launch = now - timedelta(days=10)

    async def aged(name: str, days: float) -> Conversation:
        conv = await deep_chat(world, name)

        async def apply(row: Conversation) -> None:
            row.last_activity_at = now - timedelta(days=days)

        return await world.conversations._update(conv, apply)  # noqa: SLF001

    old = await aged("old", 30)
    a = await aged("a", 5)
    b = await aged("b", 4)
    c = await aged("c", 3)
    d = await aged("d", 2.5)
    fresh = await aged("fresh", 1)
    no_session = await world.conversations.create(kind="deep", agent="d", origin="t")

    idle = await world.conversations.idle(kind="deep", days=2, since=launch)
    assert [x.external_id for x in idle] == [
        a.external_id,
        b.external_id,
        c.external_id,
        d.external_id,
    ]
    assert old.external_id not in {x.external_id for x in idle}
    assert fresh.external_id not in {x.external_id for x in idle}
    assert no_session.external_id not in {x.external_id for x in idle}

    scheduler = Scheduler(conversations=world.conversations)
    run = JobRun(Job("закрытие", lambda _run: asyncio.sleep(0)), "cron", {}, scheduler)
    closed = await run.close_idle(kind="deep", days=2, limit=3, since=launch)
    assert [r.conversation.external_id for r in closed] == [
        a.external_id,
        b.external_id,
        c.external_id,
    ]
    assert (await world.conversations.get(d.external_id)).status == "open"
    assert len(await world.conversations.idle(kind="deep", days=2, since=launch)) == 1


async def test_line_cap_bounces_a_long_rewrite_and_asks_to_shorten(
    world: World,
) -> None:
    distiller(world, CapClient)
    state = world.root / "vault" / "состояние.md"
    state.write_text("# состояние\n- было так\n", encoding="utf-8")
    CapClient.target = state
    CapClient.lines = 70
    scheduler = Scheduler(conversations=world.conversations)
    run = JobRun(Job("память", lambda _run: asyncio.sleep(0)), "cron", {}, scheduler)

    job = await run.spawn_job(
        agent="x", text="перепиши", line_cap=LineCap(state, max_lines=60)
    )
    await world.settle(job, 2)

    client = ScriptedClient.instances[-1]
    assert len(client.prompts) == 2
    assert "70 строк при потолке 60" in client.prompts[1]
    assert "[инжект: потолок" in client.prompts[1]
    assert state.read_text(encoding="utf-8").count("\n") == 10 - 1
    row = await world.conversations.get(job.external_id)
    assert row.flags["line_cap_attempts"] == 1