from __future__ import annotations

import logging
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import UTC, datetime
from threading import RLock
from typing import Any, Literal, Protocol
from uuid import uuid4

from fastapi import APIRouter
from missioncore_plugin_sdk.v0alpha2 import (
    RuntimeActionInvocation,
    RuntimeActionResult,
    RuntimeHandshakeRequest,
    RuntimeHandshakeResult,
    RuntimeHealthSnapshot,
    RuntimePluginDescriptor,
)
from pydantic import BaseModel, ConfigDict, Field, ValidationError

from k1link.sessions.plugin_contract import ObservationRuntimeContribution

STATE_READ_ACTION_ID = "state.read"

logger = logging.getLogger(__name__)


class DevicePluginActionRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    input: dict[str, Any] = Field(default_factory=dict)


class DevicePluginActionAdapter(Protocol):
    plugin_id: str
    action_ids: frozenset[str]

    async def invoke(self, invocation: RuntimeActionInvocation) -> Mapping[str, Any]: ...


def _noop() -> None:
    return None


class DevicePluginRuntimeTransport(Protocol):
    """Replaceable control-plane seam between Core and one plugin runtime."""

    descriptor: RuntimePluginDescriptor

    def handshake(self, request: RuntimeHandshakeRequest) -> RuntimeHandshakeResult: ...

    def health(self) -> RuntimeHealthSnapshot: ...

    async def invoke(self, invocation: RuntimeActionInvocation) -> RuntimeActionResult: ...

    def close(self) -> None: ...


class PluginNotFoundError(LookupError):
    """The requested plugin is not installed in the runtime composition."""


class PluginActionNotFoundError(LookupError):
    """The requested action is not declared by the selected plugin."""


class PluginExecutionError(RuntimeError):
    """A validated plugin action failed while talking to its device/runtime."""


class PluginRuntimeCompatibilityError(RuntimeError):
    """A plugin runtime cannot satisfy the reviewed manifest/host contract."""


class PluginRuntimeUnavailableError(RuntimeError):
    """A plugin runtime has not completed activation or has already stopped."""


class InProcessDevicePluginRuntime:
    """Laboratory transport preserving the future process boundary in-process."""

    def __init__(
        self,
        adapter: DevicePluginActionAdapter,
        descriptor: RuntimePluginDescriptor,
        *,
        close: Callable[[], None] = _noop,
    ) -> None:
        if adapter.plugin_id != descriptor.plugin_id:
            raise PluginRuntimeCompatibilityError(
                "In-process adapter/runtime descriptor plugin id mismatch"
            )
        if adapter.action_ids != frozenset(descriptor.action_ids):
            raise PluginRuntimeCompatibilityError(
                "In-process adapter/runtime descriptor actions mismatch"
            )
        self._adapter = adapter
        self.descriptor = descriptor
        self._close = close
        self._runtime_instance_id = f"{descriptor.plugin_id}:{uuid4().hex}"
        self._status: Literal["starting", "ready", "stopped"] = "starting"
        self._lock = RLock()

    def handshake(self, request: RuntimeHandshakeRequest) -> RuntimeHandshakeResult:
        with self._lock:
            if self._status == "stopped":
                raise PluginRuntimeUnavailableError(
                    f"Device plugin runtime is stopped: {self.descriptor.plugin_id}"
                )
            if request.plugin_id != self.descriptor.plugin_id:
                raise PluginRuntimeCompatibilityError(
                    "Runtime handshake plugin id does not match the descriptor"
                )
            if request.plugin_version != self.descriptor.plugin_version:
                raise PluginRuntimeCompatibilityError(
                    "Runtime handshake plugin version does not match the descriptor"
                )
            if request.runtime_protocol_version != self.descriptor.runtime_protocol_version:
                raise PluginRuntimeCompatibilityError(
                    "Runtime handshake protocol version is not supported"
                )
            if request.host_api_version not in self.descriptor.supported_host_api_versions:
                raise PluginRuntimeCompatibilityError(
                    "Runtime does not support the manifest host API version"
                )
            if frozenset(request.required_action_ids) != frozenset(
                self.descriptor.action_ids
            ):
                raise PluginRuntimeCompatibilityError(
                    "Runtime handshake actions do not match the manifest"
                )
            self._status = "ready"
            return RuntimeHandshakeResult(
                handshake_id=request.handshake_id,
                runtime_instance_id=self._runtime_instance_id,
                accepted_host_api_version=request.host_api_version,
                descriptor=self.descriptor,
                ready_at=datetime.now(UTC),
            )

    def health(self) -> RuntimeHealthSnapshot:
        with self._lock:
            status = self._status
        return RuntimeHealthSnapshot(
            runtime_instance_id=self._runtime_instance_id,
            plugin_id=self.descriptor.plugin_id,
            plugin_version=self.descriptor.plugin_version,
            runtime_protocol_version=self.descriptor.runtime_protocol_version,
            status=status,
            observed_at=datetime.now(UTC),
        )

    async def invoke(self, invocation: RuntimeActionInvocation) -> RuntimeActionResult:
        with self._lock:
            if self._status != "ready":
                raise PluginRuntimeUnavailableError(
                    f"Device plugin runtime is not ready: {self.descriptor.plugin_id}"
                )
        if invocation.plugin_id != self.descriptor.plugin_id:
            raise PluginActionNotFoundError("Runtime invocation targets another plugin")
        if invocation.action_id not in self.descriptor.action_ids:
            raise PluginActionNotFoundError(
                f"Device plugin {self.descriptor.plugin_id} does not declare action "
                f"{invocation.action_id}"
            )
        output = await self._adapter.invoke(invocation)
        try:
            return RuntimeActionResult(
                invocation_id=invocation.invocation_id,
                plugin_id=invocation.plugin_id,
                action_id=invocation.action_id,
                completed_at=datetime.now(UTC),
                output=dict(output),
            )
        except ValidationError as exc:
            # Never allow a malformed plugin result to masquerade as invalid
            # operator input (HTTP 422). Do not log the output or validation
            # input: device snapshots may contain retained evidence.
            logger.error(
                "device plugin returned a non-JSON action result: "
                "plugin_id=%s action_id=%s validation_errors=%d",
                invocation.plugin_id,
                invocation.action_id,
                exc.error_count(),
            )
            raise PluginExecutionError(
                "Device plugin returned an invalid non-JSON action result"
            ) from exc

    def close(self) -> None:
        with self._lock:
            if self._status == "stopped":
                return
            self._status = "stopped"
        self._close()


@dataclass(frozen=True)
class DevicePluginRuntimeContribution:
    """One reviewed backend plugin contribution loaded from its manifest factory."""

    runtime: DevicePluginRuntimeTransport
    legacy_routers: tuple[APIRouter, ...] = ()
    observation: ObservationRuntimeContribution | None = None


class DevicePluginDispatcher:
    """Host-owned dispatcher for allowlisted, namespaced plugin actions."""

    def __init__(self, runtimes: list[DevicePluginRuntimeTransport]) -> None:
        self._runtimes: dict[str, DevicePluginRuntimeTransport] = {}
        for runtime in runtimes:
            plugin_id = runtime.descriptor.plugin_id
            if plugin_id in self._runtimes:
                raise ValueError(f"Duplicate runtime device-plugin id: {plugin_id}")
            self._runtimes[plugin_id] = runtime

    @property
    def action_declarations(self) -> dict[str, frozenset[str]]:
        return {
            plugin_id: frozenset(runtime.descriptor.action_ids)
            for plugin_id, runtime in self._runtimes.items()
        }

    @property
    def health_snapshots(self) -> dict[str, RuntimeHealthSnapshot]:
        return {
            plugin_id: runtime.health() for plugin_id, runtime in self._runtimes.items()
        }

    async def invoke(
        self,
        plugin_id: str,
        action_id: str,
        payload: Mapping[str, Any],
    ) -> dict[str, Any]:
        runtime = self._runtimes.get(plugin_id)
        if runtime is None:
            raise PluginNotFoundError(f"Device plugin is not installed: {plugin_id}")
        return await invoke_device_plugin_runtime(runtime, action_id, payload)


async def invoke_device_plugin_runtime(
    runtime: DevicePluginRuntimeTransport,
    action_id: str,
    payload: Mapping[str, Any],
) -> dict[str, Any]:
    """Validate SDK request/result envelopes around one plugin action call."""

    descriptor = runtime.descriptor
    if action_id not in descriptor.action_ids:
        raise PluginActionNotFoundError(
            f"Device plugin {descriptor.plugin_id} does not declare action {action_id}"
        )
    invocation = RuntimeActionInvocation(
        invocation_id=uuid4().hex,
        plugin_id=descriptor.plugin_id,
        action_id=action_id,
        requested_at=datetime.now(UTC),
        parameters=dict(payload),
    )
    result = await runtime.invoke(invocation)
    if (
        result.invocation_id != invocation.invocation_id
        or result.plugin_id != invocation.plugin_id
        or result.action_id != invocation.action_id
    ):
        raise PluginExecutionError("Plugin runtime returned an uncorrelated action result")
    return dict(result.output)
