"""Test-only protocol peer. This is NOT a Python replacement for the Rust host.

It exercises the shipped subprocesses without Rust/provider availability. Rust
storage, worker supervision and SessionRuntime are tested separately by Cargo.
"""
from __future__ import annotations

import asyncio
import contextlib
import json
import os
import shutil
import sys
import tempfile
import uuid
from pathlib import Path
from typing import Awaitable, Callable

from klbr_runtime.release import fingerprint
from klbr_runtime.wire import Wire

ROOT = Path(__file__).resolve().parents[3]
RUNNER = ROOT / "python/klbr-runtime/src/klbr_runtime/__main__.py"
Service = Callable[[dict], Awaitable[dict]]


def behavior_copy(target: Path, delivery: str | None = None) -> Path:
    shutil.copytree(ROOT / "behavior", target)
    if delivery is not None:
        with (target / "klbr_hooks/defaults.py").open("a") as out:
            out.write("\n" + delivery + "\n")
    return target


class Peer:
    def __init__(self, role: str = "workbench", behavior: Path | None = None,
                 skills: Path | None = None):
        self.role, self.behavior, self.skills = role, behavior, skills
        self.session, self.generation = "test-session", str(uuid.uuid4())
        self.token = str(uuid.uuid4())
        self.directory = tempfile.TemporaryDirectory(prefix="klbr-test-")
        self.process = None
        self.wire = None
        self.server = None
        self.drains = []
        self.raw = {"stdout": bytearray(), "stderr": bytearray()}
        self.requests = []

    async def start(self) -> Peer:
        loop = asyncio.get_running_loop()
        accepted = loop.create_future()
        def accept(reader, writer):
            if accepted.done():
                writer.close()
            else:
                accepted.set_result((reader, writer))
        socket_path = str(Path(self.directory.name) / "control.sock")
        self.server = await asyncio.start_unix_server(accept, socket_path)
        args = [sys.executable, "-I", "-B", "-u", str(RUNNER), "--socket", socket_path,
                "--role", self.role, "--session", self.session, "--generation", self.generation,
                "--token", self.token]
        revision = None
        if self.behavior is not None:
            revision = fingerprint(self.behavior)
            args += ["--behavior", str(self.behavior), "--revision", revision]
        skills_revision = None
        if self.skills is not None:
            skills_revision = json.loads(Path(self.skills).read_text())["revision"]
            args += ["--skills", str(self.skills)]
        environment = {key: os.environ[key] for key in
            ("PATH", "HOME", "LANG", "LC_ALL", "TMPDIR", "LD_LIBRARY_PATH", "DYLD_LIBRARY_PATH")
            if key in os.environ}
        self.process = await asyncio.create_subprocess_exec(*args, env=environment, stdin=asyncio.subprocess.DEVNULL,
                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
        async def drain(name, stream):
            while chunk := await stream.read(4096):
                self.raw[name] += chunk
                del self.raw[name][:-8192]
        self.drains = [asyncio.create_task(drain("stdout", self.process.stdout)),
                       asyncio.create_task(drain("stderr", self.process.stderr))]
        exited = asyncio.create_task(self.process.wait())
        try:
            done, _ = await asyncio.wait([accepted, exited], timeout=3, return_when=asyncio.FIRST_COMPLETED)
            if accepted not in done:
                raise RuntimeError("worker failed before hello: " + self.raw["stderr"].decode(errors="replace"))
            reader, writer = accepted.result()
            self.wire = Wire(reader, writer, self.session, self.generation)
            hello = await asyncio.wait_for(self.wire.receive(), 3)
            assert hello == {"kind":"hello", "token":self.token, "role":self.role,
                             "revision":revision,"slots":["attention.plan","delivery.review"] if revision else [],
                             "skills":skills_revision}, hello
        except BaseException:
            await self.close()
            raise
        finally:
            exited.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await exited
        return self

    async def close(self) -> None:
        if self.wire:
            self.wire.writer.close()
            with contextlib.suppress(ConnectionError):
                await self.wire.writer.wait_closed()
        if self.process and self.process.returncode is None:
            self.process.kill()
            await self.process.wait()
        if self.server:
            self.server.close()
            await self.server.wait_closed()
        for task in self.drains:
            with contextlib.suppress(asyncio.CancelledError):
                await task
        self.directory.cleanup()

    async def __aenter__(self):
        return await self.start()

    async def __aexit__(self, *_):
        await self.close()

    async def begin(self, code: str) -> str:
        operation_id = str(uuid.uuid4())
        await self.wire.send({"kind":"execute", "id":operation_id, "code":code})
        return operation_id

    async def finish(self, operation_id: str, service: Service | None = None) -> dict:
        while True:
            message = await self.wire.receive()
            if message["kind"] == "completed":
                assert message["id"] == operation_id
                return message["outcome"]
            assert message["kind"] == "host_request", message
            assert message["execution_id"] == operation_id
            self.requests.append(message)
            if service is None:
                raise AssertionError(f"unexpected host request: {message}")
            reply = await service(message)
            await self.wire.send({"kind":"host_reply", "id":message["id"],
                                 "execution_id":operation_id, "reply":reply})

    async def execute(self, code: str, service: Service | None = None, timeout: float = 3) -> dict:
        operation_id = await self.begin(code)
        return await asyncio.wait_for(self.finish(operation_id, service), timeout)

    async def invoke(self, slot: str, value: dict, timeout: float = 3) -> dict:
        operation_id = str(uuid.uuid4())
        await self.wire.send({"kind":"invoke","id":operation_id,"slot":slot,"input":value})
        return await asyncio.wait_for(self.finish(operation_id), timeout)


def output(outcome: dict, stream: str | None = None) -> str:
    return "".join(record["text"] for record in outcome.get("output", []) if stream is None or record["stream"] == stream)