from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Awaitable, Callable, Protocol

from copilot import CopilotClient, ProviderConfig, RuntimeConnection
from copilot.rpc import PermissionDecisionReject


JsonObject = dict[str, Any]
Emit = Callable[[JsonObject], Awaitable[None]]


class Session(Protocol):
    session_id: str

    def on(self, handler: Callable[[Any], None]) -> Callable[[], None]: ...

    async def send(self, prompt: str) -> str: ...

    async def abort(self) -> None: ...

    async def disconnect(self) -> None: ...


class Client(Protocol):
    async def start(self) -> None: ...

    async def stop(self) -> None: ...

    async def create_session(self, **kwargs: Any) -> Session: ...

    async def resume_session(self, session_id: str, **kwargs: Any) -> Session: ...

    async def delete_session(self, session_id: str) -> None: ...


@dataclass(frozen=True)
class SidecarConfig:
    base_url: str
    model: str
    wire_api: str
    base_directory: str
    api_key: str | None = None
    idle_timeout_seconds: int = 900
    max_sessions: int = 64

    @classmethod
    def from_environment(cls) -> SidecarConfig:
        base_url = os.environ.get("LITELLM_BASE_URL", "").strip()
        model = os.environ.get("LITELLM_MODEL", "").strip()
        wire_api = os.environ.get("LITELLM_WIRE_API", "").strip()
        base_directory = os.environ.get("COPILOT_SIDECAR_HOME", "").strip()
        missing = [
            name
            for name, value in (
                ("LITELLM_BASE_URL", base_url),
                ("LITELLM_MODEL", model),
                ("LITELLM_WIRE_API", wire_api),
                ("COPILOT_SIDECAR_HOME", base_directory),
            )
            if not value
        ]
        if missing:
            raise ValueError(f"missing required environment: {', '.join(missing)}")
        if wire_api not in ("completions", "responses"):
            raise ValueError("LITELLM_WIRE_API must be 'completions' or 'responses'")
        return cls(
            base_url=base_url,
            model=model,
            wire_api=wire_api,
            base_directory=os.path.abspath(base_directory),
            api_key=os.environ.get("LITELLM_API_KEY") or None,
            idle_timeout_seconds=max(
                1,
                int(os.environ.get("COPILOT_SESSION_IDLE_SECONDS", "900")),
            ),
            max_sessions=max(
                1,
                int(os.environ.get("COPILOT_MAX_SESSIONS", "64")),
            ),
        )


@dataclass
class Turn:
    request_id: str
    done: bool = False


@dataclass
class Conversation:
    session: Session
    unsubscribe: Callable[[], None]
    active: Turn | None = None
    last_used: float = field(default_factory=time.monotonic)


@dataclass
class ConversationGate:
    lock: asyncio.Lock
    users: int = 0


def deny_permission(*_args: Any, **_kwargs: Any) -> PermissionDecisionReject:
    return PermissionDecisionReject(feedback="The inference sidecar denies all permissions.")


class Sidecar:
    def __init__(self, client: Client, config: SidecarConfig, emit: Emit):
        self._client = client
        self._config = config
        self._emit = emit
        self._conversations: dict[str, Conversation] = {}
        self._conversations_lock = asyncio.Lock()
        self._conversation_gates: dict[str, ConversationGate] = {}
        self._conversation_gates_lock = asyncio.Lock()
        self._dispatch_condition = asyncio.Condition()
        self._active_dispatches = 0
        self._event_tasks: set[asyncio.Task[None]] = set()
        self._cleanup_task: asyncio.Task[None] | None = None
        self._started = False
        self.shutting_down = False

    async def start(self) -> None:
        if not self._started:
            await self._client.start()
            self._started = True
            self._cleanup_task = asyncio.create_task(self._cleanup_loop())

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

    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:
            normalized_request_id = request_id if isinstance(request_id, str) else None
            normalized_conversation_id = (
                conversation_id if isinstance(conversation_id, str) else None
            )
            await self._send(
                "turn.error",
                normalized_request_id,
                normalized_conversation_id,
                error={"code": "invalid_request", "message": "request_id is required"},
            )
            await self._send(
                "turn.done",
                normalized_request_id,
                normalized_conversation_id,
                failed=True,
            )
            return

        if command_name == "shutdown":
            await self._shutdown(request_id, conversation_id)
            return

        if not await self._begin_dispatch():
            await self._fail(
                request_id,
                conversation_id if isinstance(conversation_id, str) else None,
                "shutting_down",
                "the sidecar is shutting down",
            )
            return
        try:
            if command_name == "health":
                await self._send(
                    "ready",
                    request_id,
                    conversation_id if isinstance(conversation_id, str) else None,
                    status="ok",
                )
                return
            if not isinstance(conversation_id, str) or not conversation_id:
                await self._fail(
                    request_id, None, "invalid_request", "conversation_id is required"
                )
                return

            await self.evict_idle_sessions(
                exclude=conversation_id,
                reserve=1 if command_name == "turn.start" else 0,
            )
            async with self._conversation_command_lock(conversation_id):
                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
                    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}",
                    )
        finally:
            await self._end_dispatch()

    @asynccontextmanager
    async def _conversation_command_lock(
        self, conversation_id: str
    ) -> AsyncIterator[None]:
        async with self._conversation_gates_lock:
            gate = self._conversation_gates.get(conversation_id)
            if gate is None:
                gate = ConversationGate(lock=asyncio.Lock())
                self._conversation_gates[conversation_id] = gate
            gate.users += 1
        try:
            async with gate.lock:
                yield
        finally:
            async with self._conversation_gates_lock:
                gate.users -= 1
                if (
                    gate.users == 0
                    and self._conversation_gates.get(conversation_id) is gate
                ):
                    del self._conversation_gates[conversation_id]

    async def _begin_dispatch(self) -> bool:
        async with self._dispatch_condition:
            if self.shutting_down:
                return False
            self._active_dispatches += 1
            return True

    async def _end_dispatch(self) -> None:
        async with self._dispatch_condition:
            self._active_dispatches -= 1
            if self._active_dispatches == 0:
                self._dispatch_condition.notify_all()

    async def _session_options(self) -> JsonObject:
        provider: ProviderConfig = {
            "type": "openai",
            "base_url": self._config.base_url,
            "wire_api": self._config.wire_api,
        }
        if self._config.api_key is not None:
            provider["api_key"] = self._config.api_key
        return {
            "on_permission_request": deny_permission,
            "model": self._config.model,
            "provider": provider,
            "streaming": True,
            "tools": [],
            "available_tools": [],
            "mcp_servers": {},
            "enable_config_discovery": False,
            "skip_custom_instructions": True,
            "enable_skills": False,
            "enable_session_store": True,
        }

    async def _get_conversation(self, conversation_id: str) -> Conversation:
        async with self._conversations_lock:
            existing = self._conversations.get(conversation_id)
            if existing is not None:
                existing.last_used = time.monotonic()
                return existing
            if len(self._conversations) >= self._config.max_sessions:
                raise RuntimeError("Copilot session capacity exhausted")

            options = await self._session_options()
            try:
                session = await self._client.resume_session(conversation_id, **options)
            except Exception:
                session = await self._client.create_session(
                    session_id=conversation_id, **options
                )
            if session is None:
                session = await self._client.create_session(
                    session_id=conversation_id, **options
                )

            def handle_event(event: Any) -> None:
                task = asyncio.create_task(self._handle_event(conversation_id, event))
                self._event_tasks.add(task)
                task.add_done_callback(self._event_tasks.discard)

            unsubscribe = session.on(handle_event)
            conversation = Conversation(session=session, unsubscribe=unsubscribe)
            self._conversations[conversation_id] = conversation
            return conversation

    async def _cleanup_loop(self) -> None:
        interval = min(60, self._config.idle_timeout_seconds)
        try:
            while not self.shutting_down:
                await asyncio.sleep(interval)
                try:
                    await self.evict_idle_sessions()
                except Exception as error:
                    print(
                        f"copilot-sidecar: session cleanup failed: {error}",
                        file=sys.stderr,
                    )
        except asyncio.CancelledError:
            return

    async def evict_idle_sessions(
        self,
        exclude: str | None = None,
        reserve: int = 0,
    ) -> None:
        now = time.monotonic()
        async with self._conversations_lock:
            inactive = sorted(
                (
                    (conversation.last_used, conversation_id)
                    for conversation_id, conversation
                    in self._conversations.items()
                    if conversation.active is None
                    and conversation_id != exclude
                )
            )
            overflow = max(
                0,
                len(self._conversations) + reserve -
                self._config.max_sessions,
            )
            candidates = {
                conversation_id
                for index, (last_used, conversation_id) in enumerate(inactive)
                if index < overflow
                or now - last_used >= self._config.idle_timeout_seconds
            }

        for conversation_id in candidates:
            async with self._conversation_command_lock(conversation_id):
                async with self._conversations_lock:
                    conversation = self._conversations.get(conversation_id)
                    if conversation is None or conversation.active is not None:
                        continue
                    over_limit = (
                        len(self._conversations) + reserve >
                        self._config.max_sessions
                    )
                    expired = (
                        time.monotonic() - conversation.last_used
                        >= self._config.idle_timeout_seconds
                    )
                    if not over_limit and not expired:
                        continue
                    self._conversations.pop(conversation_id, None)
                try:
                    conversation.unsubscribe()
                except Exception as error:
                    print(
                        f"copilot-sidecar: unsubscribe failed: {error}",
                        file=sys.stderr,
                    )
                try:
                    await conversation.session.disconnect()
                except Exception as error:
                    print(
                        f"copilot-sidecar: disconnect failed: {error}",
                        file=sys.stderr,
                    )

    async def _start_turn(
        self, request_id: str, conversation_id: str, prompt: str
    ) -> None:
        try:
            conversation = await self._get_conversation(conversation_id)
            if conversation.active is not None and not conversation.active.done:
                await self._fail(
                    request_id,
                    conversation_id,
                    "turn_in_progress",
                    "the conversation already has an active turn",
                )
                return
            conversation.active = Turn(request_id=request_id)
            await self._send("turn.accepted", request_id, conversation_id)
            await conversation.session.send(prompt)
        except Exception as error:
            await self._finish_with_error(request_id, conversation_id, error)

    async def _abort_turn(self, request_id: str, conversation_id: str) -> None:
        conversation = self._conversations.get(conversation_id)
        if conversation is None or conversation.active is None:
            await self._fail(
                request_id, conversation_id, "no_active_turn", "no active turn to abort"
            )
            return
        try:
            target_request_id = conversation.active.request_id
            await self._send(
                "turn.accepted",
                request_id,
                conversation_id,
                action="abort",
                target_request_id=target_request_id,
            )
            await conversation.session.abort()
            await self._send(
                "turn.done",
                request_id,
                conversation_id,
                action="abort",
                target_request_id=target_request_id,
            )
        except Exception as error:
            await self._fail(request_id, conversation_id, "sdk_error", str(error))

    async def _delete_conversation(
        self, request_id: str, conversation_id: str
    ) -> None:
        try:
            async with self._conversations_lock:
                conversation = self._conversations.pop(conversation_id, None)
            if conversation is not None:
                if conversation.active is not None and not conversation.active.done:
                    await self._finish_turn(
                        conversation_id, conversation.active, aborted=True, deleted=True
                    )
                conversation.unsubscribe()
                await conversation.session.disconnect()
                session_id = conversation.session.session_id
            else:
                session_id = conversation_id
            await self._client.delete_session(session_id)
            await self._send(
                "turn.done", request_id, conversation_id, action="conversation.delete"
            )
        except Exception as error:
            await self._finish_with_error(request_id, conversation_id, error)

    async def _handle_event(self, conversation_id: str, event: Any) -> None:
        conversation = self._conversations.get(conversation_id)
        if conversation is None or conversation.active is None:
            return
        turn = conversation.active
        event_type = getattr(getattr(event, "type", None), "value", None)
        data = getattr(event, "data", None)

        if event_type == "assistant.message_delta":
            await self._send(
                "assistant.delta",
                turn.request_id,
                conversation_id,
                delta=getattr(data, "delta_content", ""),
                message_id=getattr(data, "message_id", None),
            )
        elif event_type == "assistant.message":
            await self._send(
                "assistant.completed",
                turn.request_id,
                conversation_id,
                content=getattr(data, "content", ""),
                message_id=getattr(data, "message_id", None),
                model=getattr(data, "model", None),
            )
        elif event_type == "assistant.usage":
            usage = {
                name: getattr(data, name, None)
                for name in (
                    "model",
                    "input_tokens",
                    "output_tokens",
                    "reasoning_tokens",
                    "cache_read_tokens",
                    "cache_write_tokens",
                    "finish_reason",
                    "cost",
                )
                if getattr(data, name, None) is not None
            }
            await self._send(
                "assistant.usage",
                turn.request_id,
                conversation_id,
                usage=usage,
            )
        elif event_type == "session.error":
            error = {
                "code": getattr(data, "error_code", None)
                or getattr(data, "error_type", "sdk_error"),
                "message": getattr(data, "message", "Copilot session error"),
            }
            if getattr(data, "status_code", None) is not None:
                error["status_code"] = data.status_code
            await self._send(
                "turn.error",
                turn.request_id,
                conversation_id,
                error=error,
            )
            await self._finish_turn(conversation_id, turn, failed=True)
        elif event_type == "session.idle":
            await self._finish_turn(
                conversation_id, turn, aborted=bool(getattr(data, "aborted", False))
            )

    async def _finish_with_error(
        self, request_id: str, conversation_id: str | None, error: Exception
    ) -> None:
        await self._send(
            "turn.error",
            request_id,
            conversation_id,
            error={"code": "sdk_error", "message": str(error)},
        )
        conversation = (
            self._conversations.get(conversation_id)
            if conversation_id is not None
            else None
        )
        if conversation is not None and conversation.active is not None:
            await self._finish_turn(conversation_id, conversation.active, failed=True)
        else:
            await self._send("turn.done", request_id, conversation_id, failed=True)

    async def _fail(
        self,
        request_id: str,
        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)

    async def _finish_turn(
        self, conversation_id: str, turn: Turn, **fields: Any
    ) -> None:
        if turn.done:
            return
        turn.done = True
        await self._send("turn.done", turn.request_id, conversation_id, **fields)
        conversation = self._conversations.get(conversation_id)
        if conversation is not None and conversation.active is turn:
            conversation.active = None
            conversation.last_used = time.monotonic()

    async def _shutdown(
        self, request_id: str, conversation_id: Any
    ) -> None:
        async with self._dispatch_condition:
            if self.shutting_down:
                await self._send(
                    "turn.done",
                    request_id,
                    conversation_id if isinstance(conversation_id, str) else None,
                    action="shutdown",
                    already_in_progress=True,
                )
                return
            self.shutting_down = True
            if self._cleanup_task is not None:
                self._cleanup_task.cancel()
                try:
                    await self._cleanup_task
                except asyncio.CancelledError:
                    pass
                self._cleanup_task = None
            while self._active_dispatches:
                await self._dispatch_condition.wait()

        await self.drain_events()
        conversations = list(self._conversations.items())
        self._conversations.clear()
        shutdown_errors: list[str] = []
        for item_conversation_id, conversation in conversations:
            conversation.unsubscribe()
            try:
                if conversation.active is not None:
                    await self._finish_turn(
                        item_conversation_id,
                        conversation.active,
                        aborted=True,
                        shutdown=True,
                    )
                await conversation.session.disconnect()
            except Exception as error:
                shutdown_errors.append(str(error))
        if self._started:
            try:
                await self._client.stop()
            except Exception as error:
                shutdown_errors.append(str(error))
            self._started = False
        if shutdown_errors:
            await self._send(
                "turn.error",
                request_id,
                conversation_id if isinstance(conversation_id, str) else None,
                error={
                    "code": "shutdown_error",
                    "message": "; ".join(shutdown_errors),
                },
            )
        await self._send(
            "turn.done",
            request_id,
            conversation_id if isinstance(conversation_id, str) else None,
            action="shutdown",
            failed=bool(shutdown_errors),
        )

    async def drain_events(self) -> None:
        while self._event_tasks:
            await asyncio.gather(*tuple(self._event_tasks))

    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,
            }
        )


class NdjsonWriter:
    def __init__(self) -> None:
        self._lock = asyncio.Lock()

    async def __call__(self, payload: JsonObject) -> None:
        encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
        async with self._lock:
            sys.stdout.write(encoded + "\n")
            sys.stdout.flush()


def build_client(cli_path: str, base_directory: str) -> CopilotClient:
    resolved_path = os.path.abspath(cli_path)
    if not os.path.isfile(resolved_path):
        raise ValueError(f"Copilot CLI does not exist: {resolved_path}")
    if not os.access(resolved_path, os.X_OK):
        raise ValueError(f"Copilot CLI is not executable: {resolved_path}")
    os.makedirs(base_directory, mode=0o700, exist_ok=True)
    return CopilotClient(
        connection=RuntimeConnection.for_stdio(path=resolved_path),
        base_directory=base_directory,
        use_logged_in_user=False,
        log_level="error",
        mode="empty",
    )


async def run(cli_path: str) -> int:
    writer = NdjsonWriter()
    config = SidecarConfig.from_environment()
    sidecar = Sidecar(build_client(cli_path, config.base_directory), config, writer)
    await sidecar.start()
    await sidecar.announce_ready()

    tasks: set[asyncio.Task[None]] = set()
    try:
        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:
                for event_type in ("turn.error", "turn.done"):
                    payload: JsonObject = {
                        "type": event_type,
                        "request_id": None,
                        "conversation_id": None,
                    }
                    if event_type == "turn.error":
                        payload["error"] = {
                            "code": "invalid_json",
                            "message": str(error),
                        }
                    else:
                        payload["failed"] = True
                    await writer(payload)
                continue
            task = asyncio.create_task(sidecar.dispatch(command))
            tasks.add(task)
            task.add_done_callback(tasks.discard)
            if command.get("command") == "shutdown":
                await task
                break
        if tasks:
            await asyncio.gather(*tasks)
        await sidecar.drain_events()
    finally:
        if not sidecar.shutting_down:
            await sidecar.dispatch(
                {
                    "command": "shutdown",
                    "request_id": "stdin-eof",
                    "conversation_id": None,
                }
            )
    return 0


def main() -> None:
    parser = argparse.ArgumentParser(description="Copilot SDK NDJSON sidecar")
    parser.add_argument("copilot_cli", help="path to the Bazel-pinned Copilot CLI")
    args = parser.parse_args()
    try:
        raise SystemExit(asyncio.run(run(args.copilot_cli)))
    except (OSError, ValueError, RuntimeError) as error:
        print(f"copilot-sidecar: {error}", file=sys.stderr)
        raise SystemExit(2) from error


if __name__ == "__main__":
    main()