"""Governed tenant runtime and receipt boundary."""

from __future__ import annotations

import time
from pathlib import Path
from typing import Iterable

from .backends.base import Backend, BackendInput, BackendResult
from .backends.mock import EvidenceBackend
from .cards import CardContext, DeclarativeCard, PreparedRequest
from .contracts import (
    CardPackage,
    ExecutionReceipt,
    GenerationRequest,
    RuntimeResult,
    TokenUsage,
)
from .errors import AuthenticationError, PolicyViolation, UnsupportedCapability
from .receipts import new_id, sha256_value
from .scheduler import BatchScheduler
from .store import StateStore


class BackplaneRuntime:
    def __init__(
        self,
        state_path: str | Path,
        backend: Backend | None = None,
        *,
        max_batch_size: int = 8,
        batch_window_ms: int = 8,
    ):
        self.store = StateStore(state_path)
        self.backend = backend or EvidenceBackend()
        self.cards: dict[str, DeclarativeCard] = {}
        self.scheduler = BatchScheduler(
            self.backend,
            max_batch_size=max_batch_size,
            batch_window_ms=batch_window_ms,
        )

    def register(self, package: CardPackage) -> None:
        card_id = package.manifest.id
        if card_id in self.cards:
            raise ValueError(f"card already registered: {card_id}")
        self.cards[card_id] = DeclarativeCard(package)

    def register_many(self, packages: Iterable[CardPackage]) -> None:
        for package in packages:
            self.register(package)

    def assign(self, tenant_id: str, card_ids: list[str]) -> None:
        unknown = [card_id for card_id in card_ids if card_id not in self.cards]
        if unknown:
            raise ValueError(f"unknown cards: {', '.join(unknown)}")
        self.store.assign_cards(tenant_id, card_ids)

    def create_tenant(self, tenant_id: str, name: str) -> str:
        return self.store.create_tenant(tenant_id, name)

    async def run(
        self,
        request: GenerationRequest,
        *,
        tenant_key: str | None = None,
        authenticate: bool = True,
    ) -> RuntimeResult:
        if authenticate:
            if not tenant_key:
                raise AuthenticationError("tenant credential required")
            self.store.verify_tenant_key(request.tenant_id, tenant_key)
        elif not self.store.tenant_exists(request.tenant_id):
            raise AuthenticationError("unknown tenant")

        request_id = request.request_id or new_id("req")
        card_ids = self.store.card_ids_for(request.tenant_id)
        card_stack = [self.cards[card_id] for card_id in card_ids]
        stack_descriptor = [
            {
                "id": card.manifest.id,
                "version": card.manifest.version,
                "digest": card.package.digest_sha256,
            }
            for card in card_stack
        ]
        stack_digest = sha256_value(stack_descriptor)
        context = CardContext(
            tenant_id=request.tenant_id,
            request_id=request_id,
            memory=self.store.memory_for(request.tenant_id),
        )
        prepared = PreparedRequest(request=request, prompt=request.prompt)
        started = time.perf_counter()

        try:
            for card in card_stack:
                prepared = card.before(prepared, context)
            has_steering = any(
                spec.get("kind") == "steering"
                for spec in context.activation_specs
            )
            has_adapter = any(
                spec.get("kind") == "adapter"
                for spec in context.activation_specs
            )
            if has_steering and not self.backend.capabilities.activation_cards:
                raise UnsupportedCapability(
                    f"{self.backend.capabilities.backend} backend does not support "
                    "steering cards"
                )
            if has_adapter and not self.backend.capabilities.per_request_adapters:
                raise UnsupportedCapability(
                    f"{self.backend.capabilities.backend} backend does not support "
                    "per-request adapter cards"
                )
            if context.activation_specs:
                stack_digest = sha256_value(
                    {
                        "assigned_stack": stack_descriptor,
                        "active_specs": context.activation_specs,
                    }
                )
            if context.short_circuit is not None:
                backend_result = BackendResult(
                    text=context.short_circuit,
                    input_tokens=max(1, len(prepared.prompt.split())),
                    output_tokens=max(1, len(context.short_circuit.split())),
                    latency_ms=0.0,
                    metadata={"short_circuit": "exact-compute"},
                )
                queue_ms, batch_size = 0.0, 1
            else:
                backend_input = BackendInput(
                    prompt=prepared.prompt,
                    max_new_tokens=request.max_new_tokens,
                    tenant_id=request.tenant_id,
                    request_id=request_id,
                    stack_key=stack_digest,
                    activation_specs=tuple(context.activation_specs),
                    metadata=request.metadata,
                )
                backend_result, queue_ms, batch_size = await self.scheduler.submit(
                    backend_input
                )
            text = backend_result.text
            for card in reversed(card_stack):
                text = card.after(text, context)
            status = "completed"
        except PolicyViolation as exc:
            text = "Request blocked by the active tenant policy."
            status = "blocked"
            queue_ms, batch_size = 0.0, 1
            backend_result = BackendResult(
                text=text,
                input_tokens=max(1, len(request.prompt.split())),
                output_tokens=len(text.split()),
                latency_ms=0.0,
                metadata={"policy": str(exc)},
            )

        latency_ms = (time.perf_counter() - started) * 1000
        receipt = ExecutionReceipt(
            receipt_id=new_id("rcpt"),
            request_id=request_id,
            tenant_id=request.tenant_id,
            status=status,
            backend=self.backend.capabilities.backend,
            model=self.backend.capabilities.model,
            card_stack=tuple(card_ids),
            stack_digest_sha256=stack_digest,
            request_digest_sha256=sha256_value(
                {
                    "tenant": request.tenant_id,
                    "prompt": request.prompt,
                    "max_new_tokens": request.max_new_tokens,
                }
            ),
            response_digest_sha256=sha256_value({"text": text}),
            latency_ms=round(latency_ms, 3),
            queue_ms=round(queue_ms, 3),
            batch_size=batch_size,
            token_usage=TokenUsage(
                input_tokens=backend_result.input_tokens,
                output_tokens=backend_result.output_tokens,
            ),
            gpu_memory_bytes=backend_result.gpu_memory_bytes,
            events=tuple(context.events),
            limitation_notes=self.backend.capabilities.limitations,
        )
        self.store.save_receipt(
            receipt.receipt_id,
            request.tenant_id,
            receipt.model_dump(mode="json"),
        )
        return RuntimeResult(text=text, receipt=receipt)

    async def close(self) -> None:
        await self.scheduler.close()

