"""Scripted container peer implementing NDJSON contract revision 1."""

from __future__ import annotations

import asyncio
import base64
import json
from typing import Any, Callable, Dict, List, Optional


class ScriptedContainerPeer:
    """Scripted container peer for deterministic runner tests."""

    def __init__(
        self,
        scenario: str = "standard",  # "standard", "no_artifacts", "skip_artifact", "agent_error", "timeout"
        custom_steps: Optional[Callable] = None,
    ):
        self.scenario = scenario
        self.custom_steps = custom_steps
        self.received_messages: List[Dict[str, Any]] = []

    async def _read_line(self, reader: asyncio.StreamReader) -> Dict[str, Any]:
        line = await reader.readline()
        if not line:
            raise EOFError("Unexpected EOF from backend")
        msg = json.loads(line.decode("utf-8").strip())
        self.received_messages.append(msg)
        return msg

    async def _write_msg(self, writer: asyncio.StreamWriter, msg: Dict[str, Any]) -> None:
        raw = json.dumps(msg).encode("utf-8") + b"\n"
        writer.write(raw)
        await writer.drain()

    async def run(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        """Run the scripted protocol interaction."""
        if self.custom_steps:
            await self.custom_steps(reader, writer, self)
            return

        # 1. Receive 'start' frame
        start_msg = await self._read_line(reader)
        assert start_msg.get("type") == "start"

        if self.scenario == "agent_error":
            await self._write_msg(writer, {
                "v": 1,
                "type": "error",
                "id": "a_err",
                "payload": {"code": "execution_failed", "message": "Agent fatal error"},
            })
            return

        if self.scenario == "timeout":
            # Hang until cancelled or timeout
            await asyncio.sleep(300)
            return

        # 2. Tool call: confluence_search
        await self._write_msg(writer, {
            "v": 1,
            "type": "tool_request",
            "id": "a_1",
            "payload": {
                "tool": "confluence_search",
                "parameters": {"query": "deploy service X", "limit": 10},
            },
        })
        search_resp = await self._read_line(reader)
        assert search_resp.get("type") == "tool_response"
        assert search_resp.get("reply_to") == "a_1"

        # 3. Tool call: confluence_view
        await self._write_msg(writer, {
            "v": 1,
            "type": "tool_request",
            "id": "a_2",
            "payload": {
                "tool": "confluence_view",
                "parameters": {"page_id": "847291"},
            },
        })
        view_resp = await self._read_line(reader)
        assert view_resp.get("type") == "tool_response"
        assert view_resp.get("reply_to") == "a_2"

        # 4. Model call
        await self._write_msg(writer, {
            "v": 1,
            "type": "model_request",
            "id": "a_3",
            "payload": {
                "messages": [{"role": "user", "content": [{"type": "text", "text": "Deploy checklist"}]}],
                "tools": [],
            },
        })
        model_resp = await self._read_line(reader)
        assert model_resp.get("type") == "model_response"
        assert model_resp.get("reply_to") == "a_3"

        # 5. Collection start
        markdown = (
            "# Deployment Guide\n\n"
            "Refer to [Deployment Guide](https://approved.example.com/pages/viewpage.action?pageId=847291).\n"
            "Exported checklist: `checklist.md`.\n"
        )
        await self._write_msg(writer, {
            "v": 1,
            "type": "collection_start",
            "id": "a_4",
            "payload": {
                "markdown": markdown,
                "warnings": [],
            },
        })
        coll_ready = await self._read_line(reader)
        assert coll_ready.get("type") == "collection_ready"
        assert coll_ready.get("reply_to") == "a_4"

        # 6. Artifact export
        if self.scenario == "no_artifacts":
            # Complete with 0 transfers
            await self._write_msg(writer, {
                "v": 1,
                "type": "complete",
                "id": "a_comp",
                "payload": {"accepted_transfer_count": 0},
            })
            return

        if self.scenario == "skip_artifact":
            # Declare oversized file
            await self._write_msg(writer, {
                "v": 1,
                "type": "artifact_begin",
                "id": "a_art_beg",
                "payload": {
                    "transfer_id": "t_large",
                    "name": "large.bin",
                    "size_bytes": 15 * 1024 * 1024,
                },
            })
            ack = await self._read_line(reader)
            assert ack.get("type") == "artifact_ack"
            assert ack.get("payload", {}).get("decision") == "skip"

            # Complete with 0 accepted transfers
            await self._write_msg(writer, {
                "v": 1,
                "type": "complete",
                "id": "a_comp",
                "payload": {"accepted_transfer_count": 0},
            })
            return

        # Standard artifact transfer
        file_bytes = b"# Checklist\n\n- Deploy service X\n"
        b64_data = base64.b64encode(file_bytes).decode("utf-8")

        await self._write_msg(writer, {
            "v": 1,
            "type": "artifact_begin",
            "id": "a_5",
            "payload": {
                "transfer_id": "t_1",
                "name": "checklist.md",
                "size_bytes": len(file_bytes),
            },
        })
        ack_begin = await self._read_line(reader)
        assert ack_begin.get("type") == "artifact_ack"
        assert ack_begin.get("payload", {}).get("decision") == "accept"

        await self._write_msg(writer, {
            "v": 1,
            "type": "artifact_chunk",
            "id": "a_6",
            "payload": {
                "transfer_id": "t_1",
                "index": 0,
                "data_base64": b64_data,
            },
        })

        await self._write_msg(writer, {
            "v": 1,
            "type": "artifact_end",
            "id": "a_7",
            "payload": {
                "transfer_id": "t_1",
                "size_bytes": len(file_bytes),
                "chunks": 1,
            },
        })
        ack_end = await self._read_line(reader)
        assert ack_end.get("type") == "artifact_ack"
        assert ack_end.get("payload", {}).get("decision") == "stored"

        # 7. Complete
        await self._write_msg(writer, {
            "v": 1,
            "type": "complete",
            "id": "a_8",
            "payload": {"accepted_transfer_count": 1},
        })