"""Small, model-agnostic primitives for the unified Nelloz agent.

The existing conductor remains the transport and downstream integration layer.
This module owns the user-facing agent state: model roles, sessions, tasks,
intent classification, context shaping, and guarded tool execution.
"""

from __future__ import annotations

import asyncio
import inspect
import os
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Literal, Protocol
from uuid import uuid4


ModelRole = Literal["voice", "fast", "agent", "vision", "max", "image", "video"]
TaskStatus = Literal["running", "paused", "completed", "cancelled"]
IntentKind = Literal[
    "CHAT",
    "NEW_TASK",
    "MODIFY_TASK",
    "CANCEL_TASK",
    "PAUSE_TASK",
    "RESUME_TASK",
    "TASK_QUESTION",
    "SIDE_QUESTION",
    "GENERATE_IMAGE",
    "GENERATE_VIDEO",
]


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _id(prefix: str) -> str:
    return f"{prefix}_{uuid4().hex}"


@dataclass(frozen=True, slots=True)
class ModelCapabilities:
    text: bool = True
    vision: bool = False
    audio: bool = False
    video: bool = False
    tools: bool = False
    reasoning: bool = False


@dataclass(frozen=True, slots=True)
class ModelRequest:
    role: ModelRole
    prompt: str
    context: dict[str, Any] = field(default_factory=dict)
    options: dict[str, Any] = field(default_factory=dict)
    session_id: str | None = None
    task_id: str | None = None
    task_version: int | None = None
    signal: asyncio.Event | None = None


@dataclass(frozen=True, slots=True)
class ModelEvent:
    type: Literal["token", "audio", "progress", "completed", "error"]
    data: Any = None


@dataclass(frozen=True, slots=True)
class ModelResponse:
    text: str = ""
    provider: str = ""
    model: str = ""
    usage: dict[str, Any] = field(default_factory=dict)
    raw: Any = None


class ModelProvider(Protocol):
    provider_id: str
    local: bool
    capabilities: ModelCapabilities

    async def run(self, request: ModelRequest) -> ModelResponse: ...

    def stream(self, request: ModelRequest) -> Any: ...


@dataclass(slots=True)
class ModelBinding:
    model: str
    provider_id: str
    local: bool = False
    fallback: str | None = None


class ModelProviderError(RuntimeError):
    pass


class NoModelProvider(ModelProviderError):
    pass


class ModelRegistry:
    """Central role-to-model map. Core code never needs provider names."""

    DEFAULTS: dict[ModelRole, ModelBinding] = {
        "voice": ModelBinding("MiniCPM-o-4.5", "minicpm", local=True),
        "fast": ModelBinding("fast", "cloud-fast"),
        "agent": ModelBinding("Qwen3.8-27B", "qwen", local=True, fallback="cloud-agent"),
        "vision": ModelBinding("MiniCPM-o-4.5", "minicpm", local=True, fallback="cloud-vision"),
        "max": ModelBinding("frontier", "cloud-frontier"),
        "image": ModelBinding("Hunyuan-Image-3", "hunyuan"),
        "video": ModelBinding("MiniMax-H3", "minimax-h3"),
    }

    def __init__(self, bindings: dict[ModelRole, ModelBinding] | None = None) -> None:
        self.bindings = dict(bindings or self.DEFAULTS)
        self.providers: dict[str, ModelProvider] = {}

    @classmethod
    def from_env(cls) -> "ModelRegistry":
        bindings: dict[ModelRole, ModelBinding] = {}
        for role, default in cls.DEFAULTS.items():
            key = role.upper()
            bindings[role] = ModelBinding(
                model=os.getenv(f"NELLOZ_MODEL_{key}", default.model),
                provider_id=os.getenv(f"NELLOZ_PROVIDER_{key}", default.provider_id),
                local=os.getenv(
                    f"NELLOZ_LOCAL_{key}", "1" if default.local else "0"
                ).casefold()
                in {"1", "true", "yes", "on"},
                fallback=os.getenv(f"NELLOZ_FALLBACK_{key}", default.fallback or "") or None,
            )
        return cls(bindings)

    def bind(self, role: ModelRole, binding: ModelBinding) -> None:
        self.bindings[role] = binding

    def register(self, provider: ModelProvider) -> None:
        self.providers[provider.provider_id] = provider

    def binding(self, role: ModelRole) -> ModelBinding:
        return self.bindings[role]

    def describe(self) -> dict[str, dict[str, Any]]:
        return {
            role: {
                "model": binding.model,
                "provider": binding.provider_id,
                "local": binding.local,
                "fallback": binding.fallback,
                "registered": binding.provider_id in self.providers,
            }
            for role, binding in self.bindings.items()
        }

    def provider_for(self, role: ModelRole, *, mode: Literal["auto", "prefer_local", "cloud", "local_only"] = "auto") -> tuple[ModelBinding, ModelProvider]:
        binding = self.binding(role)
        candidates = [binding]
        if binding.fallback:
            fallback = next(
                (
                    item
                    for item in self.bindings.values()
                    if item.model == binding.fallback or item.provider_id == binding.fallback
                ),
                None,
            )
            if fallback is None and binding.fallback in self.providers:
                provider = self.providers[binding.fallback]
                fallback = ModelBinding(
                    model=binding.fallback,
                    provider_id=binding.fallback,
                    local=bool(getattr(provider, "local", False)),
                )
            if fallback is not None:
                candidates.append(fallback)
        seen: set[tuple[str, str]] = set()
        for candidate in candidates:
            if candidate is None or (candidate.model, candidate.provider_id) in seen:
                continue
            seen.add((candidate.model, candidate.provider_id))
            provider = self.providers.get(candidate.provider_id)
            if provider is None:
                continue
            is_local = bool(getattr(provider, "local", candidate.local))
            if mode == "local_only" and not is_local:
                continue
            if mode == "cloud" and is_local:
                continue
            return candidate, provider
        raise NoModelProvider(f"No provider is registered for model role {role!r}")

    async def run(
        self,
        role: ModelRole,
        request: ModelRequest,
        *,
        mode: Literal["auto", "prefer_local", "cloud", "local_only"] = "auto",
    ) -> ModelResponse:
        binding, provider = self.provider_for(role, mode=mode)
        try:
            response = await provider.run(request)
        except Exception as first_error:
            fallback_id = binding.fallback
            if not fallback_id:
                raise
            fallback = next(
                (
                    item
                    for item in self.bindings.values()
                    if item.model == fallback_id or item.provider_id == fallback_id
                ),
                None,
            )
            fallback_provider = self.providers.get(fallback.provider_id) if fallback else None
            if fallback is None or fallback_provider is None:
                raise first_error
            response = await fallback_provider.run(request)
            binding = fallback
        return ModelResponse(
            text=response.text,
            provider=response.provider or binding.provider_id,
            model=response.model or binding.model,
            usage=response.usage,
            raw=response.raw,
        )

    async def stream(
        self,
        role: ModelRole,
        request: ModelRequest,
        *,
        mode: Literal["auto", "prefer_local", "cloud", "local_only"] = "auto",
    ) -> Any:
        binding, provider = self.provider_for(role, mode=mode)
        stream_method = getattr(provider, "stream", None)
        if stream_method is None:
            response = await self.run(role, request, mode=mode)
            yield ModelEvent("completed", response)
            return
        stream = stream_method(request)
        if inspect.isawaitable(stream):
            stream = await stream
        async for event in stream:
            yield event


@dataclass(frozen=True, slots=True)
class AgentEvent:
    type: str
    session_id: str
    payload: dict[str, Any] = field(default_factory=dict)
    timestamp: datetime = field(default_factory=_now)


EventHandler = Callable[[AgentEvent], Awaitable[None] | None]


class EventBus:
    """In-process pub/sub bus. Handlers are isolated from one another."""

    def __init__(self) -> None:
        self._handlers: dict[str, list[EventHandler]] = {}

    def subscribe(self, event_type: str, handler: EventHandler) -> Callable[[], None]:
        self._handlers.setdefault(event_type, []).append(handler)

        def unsubscribe() -> None:
            handlers = self._handlers.get(event_type, [])
            if handler in handlers:
                handlers.remove(handler)

        return unsubscribe

    async def emit(self, event: AgentEvent) -> None:
        handlers = [*self._handlers.get(event.type, []), *self._handlers.get("*", [])]
        if not handlers:
            return
        for handler in handlers:
            try:
                result = handler(event)
                if inspect.isawaitable(result):
                    await result
            except Exception:
                # Event consumers are observational; a bad subscriber must not corrupt state.
                continue


@dataclass(slots=True)
class TaskState:
    id: str
    session_id: str
    goal: str
    status: TaskStatus = "running"
    version: int = 1
    constraints: dict[str, Any] = field(default_factory=dict)
    plan: Any = None
    created_at: datetime = field(default_factory=_now)
    updated_at: datetime = field(default_factory=_now)


@dataclass(slots=True)
class NellozSession:
    id: str
    conversation: list[dict[str, Any]] = field(default_factory=list)
    tasks: dict[str, TaskState] = field(default_factory=dict)
    active_task_id: str | None = None
    world: dict[str, Any] = field(default_factory=dict)
    voice: dict[str, Any] = field(default_factory=dict)
    memories: list[dict[str, Any]] = field(default_factory=list)


class TaskVersionConflict(RuntimeError):
    def __init__(self, task_id: str, expected: int, actual: int) -> None:
        super().__init__(f"Task {task_id} is at version {actual}, not {expected}")
        self.task_id = task_id
        self.expected = expected
        self.actual = actual


class SessionStore:
    def __init__(self) -> None:
        self._sessions: dict[str, NellozSession] = {}
        self._lock = asyncio.Lock()

    async def get_or_create(self, session_id: str) -> NellozSession:
        async with self._lock:
            return self._sessions.setdefault(session_id, NellozSession(id=session_id))

    async def get(self, session_id: str) -> NellozSession | None:
        async with self._lock:
            return self._sessions.get(session_id)


class TaskManager:
    def __init__(self, sessions: SessionStore, events: EventBus) -> None:
        self.sessions = sessions
        self.events = events
        self._lock = asyncio.Lock()

    async def active(self, session_id: str) -> TaskState | None:
        session = await self.sessions.get_or_create(session_id)
        task = session.tasks.get(session.active_task_id or "")
        return task if task and task.status in {"running", "paused"} else None

    async def create(self, session_id: str, goal: str, *, constraints: dict[str, Any] | None = None) -> TaskState:
        async with self._lock:
            session = await self.sessions.get_or_create(session_id)
            task = TaskState(id=_id("task"), session_id=session_id, goal=goal, constraints=dict(constraints or {}))
            session.tasks[task.id] = task
            session.active_task_id = task.id
        await self.events.emit(AgentEvent("task.created", session_id, self._dump(task)))
        return task

    async def update(
        self,
        task_id: str,
        *,
        goal: str | None = None,
        constraints: dict[str, Any] | None = None,
        plan: Any = None,
    ) -> TaskState:
        task = await self._find(task_id)
        async with self._lock:
            task.version += 1
            if goal:
                task.goal = goal
            if constraints:
                task.constraints.update(constraints)
            if plan is not None:
                task.plan = plan
            task.updated_at = _now()
        await self.events.emit(AgentEvent("task.updated", task.session_id, self._dump(task)))
        return task

    async def set_status(self, task_id: str, status: TaskStatus) -> TaskState:
        task = await self._find(task_id)
        async with self._lock:
            task.status = status
            task.version += 1
            task.updated_at = _now()
        await self.events.emit(AgentEvent(f"task.{status}", task.session_id, self._dump(task)))
        return task

    async def current(self, task_id: str, version: int) -> TaskState:
        task = await self._find(task_id)
        if task.version != version:
            raise TaskVersionConflict(task_id, version, task.version)
        if task.status != "running":
            raise TaskVersionConflict(task_id, version, task.version)
        return task

    async def _find(self, task_id: str) -> TaskState:
        # Session IDs are not encoded in task IDs, so scan the small in-memory index.
        for session_id in list(self.sessions._sessions):
            session = await self.sessions.get(session_id)
            if session and task_id in session.tasks:
                return session.tasks[task_id]
        raise KeyError(task_id)

    @staticmethod
    def _dump(task: TaskState) -> dict[str, Any]:
        return {
            "task_id": task.id,
            "goal": task.goal,
            "status": task.status,
            "version": task.version,
            "constraints": dict(task.constraints),
            "plan": task.plan,
        }


@dataclass(frozen=True, slots=True)
class Intent:
    kind: IntentKind
    text: str
    confidence: float


class IntentController:
    _STOP = re.compile(r"^(?:bitte\s+)?(?:stopp|stop|halt|abbrechen|cancel)(?:[.!?\s]*)$", re.I)
    _PAUSE = re.compile(r"^(?:bitte\s+)?(?:warte|pausier|pause)(?:[.!?\s]*)$", re.I)
    _RESUME = re.compile(r"^(?:bitte\s+)?(?:mach\s+weiter|fortsetzen|resume)(?:[.!?\s]*)$", re.I)
    _IMAGE = re.compile(r"\b(?:bild|foto|image|illustration|grafik)\b", re.I)
    _VIDEO = re.compile(r"\b(?:video|clip|film|animation)\b", re.I)
    _MODIFY = re.compile(r"\b(?:aber|stattdessen|ändere|aendere|doch|erst|nach|unter|ohne|mit)\b", re.I)

    def classify(self, text: str, active_task: TaskState | None = None) -> Intent:
        normalized = " ".join(text.split())
        if self._STOP.fullmatch(normalized):
            return Intent("CANCEL_TASK", text, 0.99)
        if self._PAUSE.fullmatch(normalized):
            return Intent("PAUSE_TASK", text, 0.99)
        if self._RESUME.fullmatch(normalized):
            return Intent("RESUME_TASK", text, 0.99)
        if self._IMAGE.search(normalized) and re.search(r"\b(?:mach|erstell|generier|create|make)\w*\b", normalized, re.I):
            return Intent("GENERATE_IMAGE", text, 0.95)
        if self._VIDEO.search(normalized) and re.search(r"\b(?:mach|erstell|generier|create|make)\w*\b", normalized, re.I):
            return Intent("GENERATE_VIDEO", text, 0.95)
        if active_task and self._MODIFY.search(normalized):
            return Intent("MODIFY_TASK", text, 0.82)
        if active_task and re.search(r"\b(?:was|welche|wie)\b", normalized, re.I):
            return Intent("TASK_QUESTION", text, 0.72)
        if re.search(r"\b(?:such|recherchier|öffne|oeffne|klick|schreib|erstelle|analysier|buche|find)\w*\b", normalized, re.I):
            return Intent("NEW_TASK", text, 0.75)
        return Intent("CHAT", text, 0.65)


class ContextBuilder:
    def __init__(self, *, recent_messages: int = 12, recent_memories: int = 8) -> None:
        self.recent_messages = recent_messages
        self.recent_memories = recent_memories

    def build(self, session: NellozSession, *, target_role: ModelRole, task: TaskState | None = None) -> dict[str, Any]:
        return {
            "target_role": target_role,
            "stable_system": "You are Nelloz, one persistent agent. Keep the user-facing identity stable.",
            "current_task": TaskManager._dump(task) if task else None,
            "world": dict(session.world),
            "voice": dict(session.voice),
            "recent_conversation": list(session.conversation[-self.recent_messages:]),
            "relevant_memories": list(session.memories[-self.recent_memories:]),
        }


@dataclass(frozen=True, slots=True)
class ToolCall:
    id: str
    session_id: str
    task_id: str | None
    task_version: int | None
    tool: str
    args: Any
    risk: Literal["read", "write", "destructive", "external"] = "read"


@dataclass(frozen=True, slots=True)
class ToolResult:
    ok: bool
    value: Any = None
    error: str | None = None


ToolHandler = Callable[[ToolCall], Awaitable[ToolResult] | ToolResult]
PermissionCheck = Callable[[ToolCall], Awaitable[bool] | bool]


class ToolBroker:
    def __init__(self, tasks: TaskManager, events: EventBus) -> None:
        self.tasks = tasks
        self.events = events
        self._handlers: dict[str, ToolHandler] = {}
        self._permissions: PermissionCheck | None = None

    def register(self, name: str, handler: ToolHandler) -> None:
        self._handlers[name] = handler

    def set_permission_check(self, check: PermissionCheck) -> None:
        self._permissions = check

    async def execute(self, call: ToolCall) -> ToolResult:
        await self.events.emit(AgentEvent("tool.requested", call.session_id, {"tool": call.tool, "task_id": call.task_id}))
        if call.task_id is not None and call.task_version is not None:
            await self.tasks.current(call.task_id, call.task_version)
        if self._permissions is not None:
            allowed = self._permissions(call)
            if inspect.isawaitable(allowed):
                allowed = await allowed
            if not allowed:
                result = ToolResult(False, error="permission_denied")
                await self.events.emit(AgentEvent("tool.failed", call.session_id, {"tool": call.tool, "error": result.error}))
                return result
        handler = self._handlers.get(call.tool)
        if handler is None:
            result = ToolResult(False, error=f"unknown_tool:{call.tool}")
            await self.events.emit(AgentEvent("tool.failed", call.session_id, {"tool": call.tool, "error": result.error}))
            return result
        await self.events.emit(AgentEvent("tool.started", call.session_id, {"tool": call.tool}))
        try:
            result = handler(call)
            result = await result if inspect.isawaitable(result) else result
            if call.task_id is not None and call.task_version is not None:
                await self.tasks.current(call.task_id, call.task_version)
        except asyncio.CancelledError:
            await self.events.emit(AgentEvent("tool.failed", call.session_id, {"tool": call.tool, "error": "cancelled"}))
            raise
        except Exception as exc:
            await self.events.emit(AgentEvent("tool.failed", call.session_id, {"tool": call.tool, "error": type(exc).__name__}))
            return ToolResult(False, error=type(exc).__name__)
        await self.events.emit(AgentEvent("tool.completed", call.session_id, {"tool": call.tool}))
        return result


class UnifiedRouter:
    def role_for(self, intent: Intent, *, needs_vision: bool = False) -> ModelRole:
        if intent.kind in {"GENERATE_IMAGE"}:
            return "image"
        if intent.kind in {"GENERATE_VIDEO"}:
            return "video"
        if needs_vision:
            return "vision"
        if intent.kind in {"NEW_TASK", "MODIFY_TASK"}:
            return "agent"
        if intent.kind in {"CHAT", "SIDE_QUESTION", "TASK_QUESTION"}:
            return "fast"
        return "voice"


@dataclass(frozen=True, slots=True)
class ModelUsageEvent:
    session_id: str
    role: ModelRole
    provider: str
    model: str
    latency_ms: float
    success: bool
    task_id: str | None = None
    usage: dict[str, Any] = field(default_factory=dict)


class Telemetry:
    def __init__(self) -> None:
        self.model_calls: list[ModelUsageEvent] = []
        self.counters: dict[str, int] = {}

    def count(self, name: str, amount: int = 1) -> None:
        self.counters[name] = self.counters.get(name, 0) + amount


class UnifiedAgentSystem:
    """Composition root for the model-agnostic agent primitives."""

    def __init__(self, *, registry: ModelRegistry | None = None) -> None:
        self.events = EventBus()
        self.models = registry or ModelRegistry.from_env()
        self.sessions = SessionStore()
        self.tasks = TaskManager(self.sessions, self.events)
        self.context = ContextBuilder()
        self.tools = ToolBroker(self.tasks, self.events)
        self.intent = IntentController()
        self.router = UnifiedRouter()
        self.telemetry = Telemetry()

    async def observe_input(self, session_id: str, text: str, *, turn_id: str) -> tuple[NellozSession, Intent]:
        session = await self.sessions.get_or_create(session_id)
        active = await self.tasks.active(session_id)
        intent = self.intent.classify(text, active)
        session.conversation.append({"role": "user", "text": text, "turn_id": turn_id, "intent": intent.kind})
        await self.events.emit(AgentEvent("user.text.final", session_id, {"text": text, "turn_id": turn_id, "intent": intent.kind}))
        return session, intent

    async def apply_intent(self, session_id: str, intent: Intent, *, route: str) -> TaskState | None:
        active = await self.tasks.active(session_id)
        if intent.kind == "MODIFY_TASK" and active is not None:
            return await self.tasks.update(active.id, constraints={"latest_user_update": intent.text}, goal=f"{active.goal}\nUser update: {intent.text}")
        if intent.kind == "PAUSE_TASK" and active is not None:
            return await self.tasks.set_status(active.id, "paused")
        if intent.kind == "RESUME_TASK" and active is not None:
            return await self.tasks.set_status(active.id, "running")
        if intent.kind == "CANCEL_TASK" and active is not None:
            return await self.tasks.set_status(active.id, "cancelled")
        if route in {"hermes", "image", "video"} and intent.kind not in {"TASK_QUESTION", "SIDE_QUESTION"}:
            return await self.tasks.create(session_id, intent.text)
        return active