from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

from pydantic import BaseModel, Field, ValidationError

from .errors import AIError
from .schemas import AccessType, RiskLevel, SupervisorPlan


class SystemStatusInput(BaseModel):
    services: list[Literal["core-api", "database", "redis", "ai", "browser"]] = Field(
        default_factory=lambda: ["core-api"], max_length=5
    )


class SystemStatusOutput(BaseModel):
    statuses: dict[str, Literal["ok", "unavailable", "unconfigured"]]


class WorkspaceProfileInput(BaseModel):
    include_permissions: bool = True


class WorkspaceProfileOutput(BaseModel):
    workspace_name: str
    timezone: str
    permissions: list[str] = Field(default_factory=list)


@dataclass(frozen=True, slots=True)
class ToolDefinition:
    name: str
    version: int
    description: str
    owning_service: str
    required_permission: str
    access_type: AccessType
    base_risk: RiskLevel
    idempotent: bool
    timeout_seconds: int
    verification_strategy: str
    input_model: type[BaseModel]
    output_model: type[BaseModel]

    def public_contract(self) -> dict:
        return {
            "name": self.name,
            "version": self.version,
            "description": self.description,
            "owning_service": self.owning_service,
            "required_permission": self.required_permission,
            "access_type": self.access_type,
            "base_risk": self.base_risk,
            "idempotent": self.idempotent,
            "timeout_seconds": self.timeout_seconds,
            "verification_strategy": self.verification_strategy,
            "input_schema": self.input_model.model_json_schema(),
            "output_schema": self.output_model.model_json_schema(),
        }


class ToolRegistry:
    def __init__(self, definitions: list[ToolDefinition] | None = None):
        self._definitions: dict[str, ToolDefinition] = {}
        for definition in definitions or []:
            self.register(definition)

    def register(self, definition: ToolDefinition) -> None:
        if definition.name in self._definitions:
            raise ValueError(f"Tool {definition.name} is already registered")
        self._definitions[definition.name] = definition

    def list(self) -> list[ToolDefinition]:
        return [self._definitions[name] for name in sorted(self._definitions)]

    def select(self, allowed: list[str]) -> list[ToolDefinition]:
        missing = sorted(set(allowed) - self._definitions.keys())
        if missing:
            raise AIError("TOOL_NOT_REGISTERED", "An allowed tool is not registered.", 422)
        return [self._definitions[name] for name in allowed]

    def validate_plan(self, plan: SupervisorPlan, allowed: list[str]) -> None:
        selected = {tool.name: tool for tool in self.select(allowed)}
        for step in plan.steps:
            tool = selected.get(step.tool_name)
            if not tool:
                raise AIError("TOOL_NOT_ALLOWED", "The plan proposed a disallowed tool.", 422)
            try:
                tool.input_model.model_validate(step.arguments)
            except ValidationError as error:
                raise AIError(
                    "TOOL_ARGUMENTS_INVALID",
                    "The plan proposed invalid tool arguments.",
                    422,
                ) from error


def default_registry() -> ToolRegistry:
    return ToolRegistry([
        ToolDefinition(
            name="system.status.read",
            version=1,
            description="Read explicit configuration and availability state for Hayva services.",
            owning_service="core-api",
            required_permission="system.status.read",
            access_type=AccessType.read,
            base_risk=RiskLevel.low,
            idempotent=True,
            timeout_seconds=5,
            verification_strategy="authoritative_health_response",
            input_model=SystemStatusInput,
            output_model=SystemStatusOutput,
        ),
        ToolDefinition(
            name="workspace.profile.read",
            version=1,
            description="Read the authorized workspace display profile and granted permissions.",
            owning_service="core-api",
            required_permission="workspace.profile.read",
            access_type=AccessType.read,
            base_risk=RiskLevel.low,
            idempotent=True,
            timeout_seconds=5,
            verification_strategy="authoritative_database_read",
            input_model=WorkspaceProfileInput,
            output_model=WorkspaceProfileOutput,
        ),
    ])