from __future__ import annotations

import argparse
import asyncio
import json
import os
import pathlib
import re
import sys
from dataclasses import dataclass
from typing import Any, Awaitable, Callable


JsonObject = dict[str, Any]
Emit = Callable[[JsonObject], Awaitable[None]]
COMMAND_PATTERN = re.compile(r"^![a-z][a-z0-9_-]*$")
EVENT_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
RESERVED_EVENT_TYPES = {"bridge.closed", "ready", "turn.accepted", "turn.done"}
RESERVED_EVENT_FIELDS = {
    "conversation_id",
    "mock",
    "mock_command",
    "request_id",
}
MAX_DELAY_MS = 60_000
PROMPT_PROFILES = {"public_visitor", "invited_friend", "june_admin"}
PROMPT_VERSION = 1
KNOWLEDGE_VERSION = 1


@dataclass(frozen=True)
class MockEvent:
    event_type: str
    delay_ms: int | None
    payload: JsonObject


@dataclass(frozen=True)
class MockCommand:
    name: str
    description: str
    events: tuple[MockEvent, ...]
    failed: bool


@dataclass(frozen=True)
class MockConfig:
    delay_ms: int
    fallback: str
    commands: dict[str, MockCommand]

    def select(self, prompt: str) -> MockCommand:
        first_token = prompt.strip().split(maxsplit=1)[0].lower()
        return self.commands.get(first_token, self.commands[self.fallback])


@dataclass
class MockTurn:
    request_id: str
    task: asyncio.Task[None]


def _require_non_negative_integer(value: Any, description: str) -> int:
    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
        raise ValueError(f"{description} must be a non-negative integer")
    if value > MAX_DELAY_MS:
        raise ValueError(f"{description} must not exceed {MAX_DELAY_MS}")
    return value


def _validate_event(command_name: str, index: int, raw: Any) -> MockEvent:
    if not isinstance(raw, dict):
        raise ValueError(f"{command_name} event {index} must be an object")
    event_type = raw.get("type")
    if not isinstance(event_type, str) or not EVENT_PATTERN.fullmatch(event_type):
        raise ValueError(
            f"{command_name} event {index} requires a safe event type"
        )
    if event_type in RESERVED_EVENT_TYPES:
        raise ValueError(
            f"{command_name} cannot script lifecycle event {event_type}"
        )
    reserved = set(raw) & RESERVED_EVENT_FIELDS
    if reserved:
        raise ValueError(
            f"{command_name} event {index} uses reserved fields: "
            f"{', '.join(sorted(reserved))}"
        )
    delay_ms = raw.get("delay_ms")
    if delay_ms is not None:
        delay_ms = _require_non_negative_integer(
            delay_ms,
            f"{command_name} event {index} delay_ms",
        )
    payload = {
        key: value
        for key, value in raw.items()
        if key not in {"delay_ms", "type"}
    }
    if event_type == "assistant.delta":
        if not isinstance(payload.get("delta"), str) or not payload["delta"]:
            raise ValueError(
                f"{command_name} assistant.delta requires non-empty delta"
            )
    elif event_type == "assistant.completed":
        if not isinstance(payload.get("content"), str):
            raise ValueError(
                f"{command_name} assistant.completed requires content"
            )
    elif event_type == "assistant.usage":
        usage = payload.get("usage")
        if not isinstance(usage, dict):
            raise ValueError(f"{command_name} assistant.usage requires usage")
        for key in ("input_tokens", "output_tokens"):
            _require_non_negative_integer(
                usage.get(key),
                f"{command_name} assistant.usage {key}",
            )
    elif event_type == "turn.error":
        error = payload.get("error")
        if not isinstance(error, dict):
            raise ValueError(f"{command_name} turn.error requires error")
        if (
            not isinstance(error.get("code"), str)
            or not error["code"]
            or not isinstance(error.get("message"), str)
            or not error["message"]
        ):
            raise ValueError(
                f"{command_name} turn.error requires error code and message"
            )
    return MockEvent(event_type, delay_ms, payload)


def _validate_command(name: str, raw: Any) -> MockCommand:
    if not COMMAND_PATTERN.fullmatch(name):
        raise ValueError(f"invalid mock command: {name}")
    if not isinstance(raw, dict):
        raise ValueError(f"{name} must contain an object")
    unknown = set(raw) - {"description", "events"}
    if unknown:
        raise ValueError(
            f"{name} has unknown keys: {', '.join(sorted(unknown))}"
        )
    description = raw.get("description", "")
    if not isinstance(description, str):
        raise ValueError(f"{name} description must be a string")
    raw_events = raw.get("events")
    if not isinstance(raw_events, list) or not raw_events:
        raise ValueError(f"{name} events must be a non-empty array")
    if len(raw_events) > 256:
        raise ValueError(f"{name} has too many events")
    events = tuple(
        _validate_event(name, index, event)
        for index, event in enumerate(raw_events)
    )
    errors = [event for event in events if event.event_type == "turn.error"]
    completed = [
        event for event in events if event.event_type == "assistant.completed"
    ]
    if errors:
        if len(errors) != 1 or errors[0] is not events[-1]:
            raise ValueError(f"{name} turn.error must be the final scripted event")
        if completed:
            raise ValueError(f"{name} cannot complete and fail the same turn")
    elif len(completed) != 1:
        raise ValueError(f"{name} requires one assistant.completed event")

    deltas = "".join(
        event.payload["delta"]
        for event in events
        if event.event_type == "assistant.delta"
    )
    if completed:
        completed_index = events.index(completed[0])
        if any(
            event.event_type == "assistant.delta"
            for event in events[completed_index + 1:]
        ):
            raise ValueError(
                f"{name} assistant.completed must follow all deltas"
            )
    if completed and deltas and deltas != completed[0].payload["content"]:
        raise ValueError(f"{name} deltas must join to completed content")
    return MockCommand(name, description, events, failed=bool(errors))


def load_mock_config(path: str | os.PathLike[str]) -> MockConfig:
    with open(path, encoding="utf-8") as config_file:
        payload = json.load(config_file)
    if not isinstance(payload, dict):
        raise ValueError("mock response file must contain a JSON object")
    unknown = set(payload) - {"commands", "delay_ms", "fallback"}
    if unknown:
        raise ValueError(f"unknown mock config keys: {', '.join(sorted(unknown))}")
    delay_ms = _require_non_negative_integer(
        payload.get("delay_ms", 55),
        "delay_ms",
    )
    fallback = payload.get("fallback")
    if not isinstance(fallback, str) or not COMMAND_PATTERN.fullmatch(fallback):
        raise ValueError("fallback must be a !command")
    raw_commands = payload.get("commands")
    if not isinstance(raw_commands, dict) or not raw_commands:
        raise ValueError("commands must be a non-empty object")
    commands = {
        name.lower(): _validate_command(name.lower(), raw)
        for name, raw in raw_commands.items()
    }
    if fallback.lower() not in commands:
        raise ValueError(f"fallback command is not defined: {fallback}")
    return MockConfig(delay_ms, fallback.lower(), commands)


class MockSidecar:
    def __init__(self, emit: Emit, config: MockConfig) -> None:
        self._emit = emit
        self._config = config
        self._active: dict[str, MockTurn] = {}
        self.shutting_down = False

    async def announce_ready(self) -> None:
        await self._send("ready", None, None, status="ok", mock=True)

    async def dispatch(self, command: JsonObject) -> None:
        command_name = command.get("command")
        request_id = command.get("request_id")
        conversation_id = command.get("conversation_id")
        if not isinstance(request_id, str) or not request_id:
            await self._fail(
                request_id if isinstance(request_id, str) else None,
                conversation_id if isinstance(conversation_id, str) else None,
                "invalid_request",
                "request_id is required",
            )
            return
        if command_name == "shutdown":
            await self._shutdown(request_id, conversation_id)
            return
        if command_name == "health":
            await self._send(
                "ready",
                request_id,
                conversation_id if isinstance(conversation_id, str) else None,
                status="ok",
                mock=True,
            )
            return
        if not isinstance(conversation_id, str) or not conversation_id:
            await self._fail(
                request_id,
                None,
                "invalid_request",
                "conversation_id is required",
            )
            return
        if command_name == "turn.start":
            prompt = command.get("prompt")
            if not isinstance(prompt, str) or not prompt:
                await self._fail(
                    request_id,
                    conversation_id,
                    "invalid_request",
                    "prompt is required",
                )
                return
            prompt_profile = command.get("prompt_profile")
            prompt_version = command.get("prompt_version")
            knowledge_version = command.get("knowledge_version")
            if (
                prompt_profile not in PROMPT_PROFILES
                or prompt_version != PROMPT_VERSION
                or knowledge_version != KNOWLEDGE_VERSION
            ):
                await self._fail(
                    request_id,
                    conversation_id,
                    "invalid_prompt_profile",
                    "a known prompt profile and current versions are required",
                )
                return
            await self._start_turn(request_id, conversation_id, prompt)
        elif command_name == "turn.abort":
            await self._abort_turn(request_id, conversation_id)
        elif command_name == "conversation.delete":
            await self._delete_conversation(request_id, conversation_id)
        else:
            await self._fail(
                request_id,
                conversation_id,
                "unknown_command",
                f"unsupported command: {command_name!r}",
            )

    async def wait_for_idle(self) -> None:
        tasks = [turn.task for turn in self._active.values()]
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    async def _start_turn(
        self,
        request_id: str,
        conversation_id: str,
        prompt: str,
    ) -> None:
        if conversation_id in self._active:
            await self._fail(
                request_id,
                conversation_id,
                "turn_in_progress",
                "the conversation already has an active turn",
            )
            return
        scripted_command = self._config.select(prompt)
        await self._send(
            "turn.accepted",
            request_id,
            conversation_id,
            mock=True,
            mock_command=scripted_command.name,
        )
        task = asyncio.create_task(
            self._stream_turn(request_id, conversation_id, scripted_command)
        )
        self._active[conversation_id] = MockTurn(request_id, task)

    async def _stream_turn(
        self,
        request_id: str,
        conversation_id: str,
        scripted_command: MockCommand,
    ) -> None:
        try:
            for event in scripted_command.events:
                delay = (
                    self._config.delay_ms
                    if event.delay_ms is None
                    else event.delay_ms
                )
                await asyncio.sleep(delay / 1000)
                await self._send(
                    event.event_type,
                    request_id,
                    conversation_id,
                    **event.payload,
                    mock=True,
                    mock_command=scripted_command.name,
                )
            await self._send(
                "turn.done",
                request_id,
                conversation_id,
                failed=scripted_command.failed,
                mock=True,
                mock_command=scripted_command.name,
            )
        except asyncio.CancelledError:
            await self._send(
                "turn.done",
                request_id,
                conversation_id,
                aborted=True,
                mock=True,
                mock_command=scripted_command.name,
            )
        except Exception as error:
            await self._send(
                "turn.error",
                request_id,
                conversation_id,
                error={
                    "code": "mock_script_failed",
                    "message": str(error),
                },
                mock_command=scripted_command.name,
            )
            await self._send(
                "turn.done",
                request_id,
                conversation_id,
                failed=True,
                mock=True,
                mock_command=scripted_command.name,
            )
        finally:
            active = self._active.get(conversation_id)
            if active is not None and active.request_id == request_id:
                self._active.pop(conversation_id, None)

    async def _abort_turn(
        self,
        request_id: str,
        conversation_id: str,
    ) -> None:
        active = self._active.get(conversation_id)
        if active is None:
            await self._fail(
                request_id,
                conversation_id,
                "no_active_turn",
                "no active turn to abort",
            )
            return
        await self._send(
            "turn.accepted",
            request_id,
            conversation_id,
            action="abort",
            target_request_id=active.request_id,
            mock=True,
        )
        active.task.cancel()
        await asyncio.gather(active.task, return_exceptions=True)
        await self._send(
            "turn.done",
            request_id,
            conversation_id,
            action="abort",
            target_request_id=active.request_id,
            mock=True,
        )

    async def _delete_conversation(
        self,
        request_id: str,
        conversation_id: str,
    ) -> None:
        active = self._active.get(conversation_id)
        if active is not None:
            active.task.cancel()
            await asyncio.gather(active.task, return_exceptions=True)
        await self._send(
            "turn.done",
            request_id,
            conversation_id,
            action="conversation.delete",
            mock=True,
        )

    async def _shutdown(
        self,
        request_id: str,
        conversation_id: Any,
    ) -> None:
        self.shutting_down = True
        tasks = [turn.task for turn in self._active.values()]
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        await self._send(
            "turn.done",
            request_id,
            conversation_id if isinstance(conversation_id, str) else None,
            action="shutdown",
            mock=True,
        )

    async def _fail(
        self,
        request_id: str | None,
        conversation_id: str | None,
        code: str,
        message: str,
    ) -> None:
        await self._send(
            "turn.error",
            request_id,
            conversation_id,
            error={"code": code, "message": message},
        )
        await self._send(
            "turn.done",
            request_id,
            conversation_id,
            failed=True,
            mock=True,
        )

    async def _send(
        self,
        event_type: str,
        request_id: str | None,
        conversation_id: str | None,
        **fields: Any,
    ) -> None:
        await self._emit(
            {
                "type": event_type,
                "request_id": request_id,
                "conversation_id": conversation_id,
                **fields,
            }
        )


async def run(config: MockConfig) -> None:
    async def emit(payload: JsonObject) -> None:
        sys.stdout.write(json.dumps(payload, separators=(",", ":")) + "\n")
        sys.stdout.flush()

    sidecar = MockSidecar(emit, config)
    await sidecar.announce_ready()
    while not sidecar.shutting_down:
        line = await asyncio.to_thread(sys.stdin.readline)
        if not line:
            break
        try:
            command = json.loads(line)
            if not isinstance(command, dict):
                raise ValueError("command must be a JSON object")
        except (json.JSONDecodeError, ValueError) as error:
            await sidecar._fail(None, None, "invalid_json", str(error))
            continue
        await sidecar.dispatch(command)
    await sidecar.wait_for_idle()


def main() -> None:
    packaged_responses = pathlib.Path(__file__).with_name("mock_responses.json")
    parser = argparse.ArgumentParser(description="Scripted JRPG mock sidecar")
    parser.add_argument("copilot_cli", nargs="?")
    parser.add_argument("--responses", default=str(packaged_responses))
    args = parser.parse_args()
    responses_path = os.environ.get("MRJUNEJUNE_MOCK_RESPONSES") or args.responses
    config = load_mock_config(responses_path)
    delay_override = os.environ.get("MRJUNEJUNE_MOCK_DELAY_MS")
    if delay_override is not None:
        try:
            delay_ms = int(delay_override)
        except ValueError as error:
            raise ValueError(
                "MRJUNEJUNE_MOCK_DELAY_MS must be a non-negative integer"
            ) from error
        delay_ms = _require_non_negative_integer(
            delay_ms,
            "MRJUNEJUNE_MOCK_DELAY_MS",
        )
        config = MockConfig(delay_ms, config.fallback, config.commands)
    asyncio.run(run(config))


if __name__ == "__main__":
    main()