from __future__ import annotations

from dataclasses import dataclass
from datetime import UTC, datetime
from importlib import import_module
from pathlib import Path
from typing import Any
from uuid import uuid4

from fastapi import APIRouter
from missioncore_plugin_sdk.v0alpha2 import RuntimeHandshakeRequest

from k1link.sessions.plugin_contract import (
    ObservationArchiveSource,
    RecordedPointColorRenderer,
    RecordingExporter,
)
from k1link.web.plugin_catalog import DevicePluginCatalog, DevicePluginManifest
from k1link.web.plugin_runtime import (
    DevicePluginDispatcher,
    DevicePluginRuntimeContribution,
)


class DevicePluginCompositionError(RuntimeError):
    """The reviewed plugin manifests and executable runtime do not match."""


@dataclass(frozen=True)
class InstalledDevicePluginEnvironment:
    catalog: DevicePluginCatalog
    dispatcher: DevicePluginDispatcher
    legacy_routers: tuple[APIRouter, ...]
    _contributions: tuple[DevicePluginRuntimeContribution, ...]

    @property
    def observation_archives(self) -> tuple[ObservationArchiveSource, ...]:
        return tuple(
            archive
            for contribution in self._contributions
            if contribution.observation is not None
            for archive in contribution.observation.archives
        )

    @property
    def recording_exporters(self) -> dict[str, RecordingExporter]:
        return {
            contribution.runtime.descriptor.plugin_id: (
                contribution.observation.recording_exporter
            )
            for contribution in self._contributions
            if contribution.observation is not None
        }

    @property
    def point_color_renderers(self) -> dict[str, RecordedPointColorRenderer]:
        return {
            contribution.runtime.descriptor.plugin_id: renderer
            for contribution in self._contributions
            if contribution.observation is not None
            if (renderer := contribution.observation.point_color_renderer) is not None
        }

    @property
    def runtime_health(self) -> tuple[dict[str, Any], ...]:
        return tuple(
            contribution.runtime.health().model_dump(mode="json")
            for contribution in self._contributions
        )

    def close(self) -> None:
        cleanup_errors = _close_contributions(self._contributions)
        if cleanup_errors:
            error = DevicePluginCompositionError(
                "One or more device plugins failed during shutdown"
            )
            _add_cleanup_notes(error, cleanup_errors)
            raise error from cleanup_errors[0]


def _close_contributions(
    contributions: tuple[DevicePluginRuntimeContribution, ...]
    | list[DevicePluginRuntimeContribution],
) -> list[Exception]:
    errors: list[Exception] = []
    for contribution in reversed(contributions):
        try:
            contribution.runtime.close()
        except Exception as exc:
            errors.append(exc)
    return errors


def _add_cleanup_notes(error: BaseException, cleanup_errors: list[Exception]) -> None:
    for index, cleanup_error in enumerate(cleanup_errors, start=1):
        error.add_note(
            f"device-plugin cleanup failure {index}: "
            f"{type(cleanup_error).__name__}: {cleanup_error}"
        )


def _load_factory(entrypoint: str) -> Any:
    module_name, separator, attribute_name = entrypoint.partition(":")
    if not separator or not module_name or not attribute_name:
        raise DevicePluginCompositionError(f"Invalid device-plugin backendEntrypoint: {entrypoint}")
    try:
        module = import_module(module_name)
        factory = getattr(module, attribute_name)
    except (ImportError, AttributeError) as exc:
        raise DevicePluginCompositionError(
            f"Cannot load device-plugin backendEntrypoint {entrypoint}: {exc}"
        ) from exc
    if not callable(factory):
        raise DevicePluginCompositionError(
            f"Device-plugin backendEntrypoint is not callable: {entrypoint}"
        )
    return factory


def _load_contribution(
    repository_root: Path,
    manifest: DevicePluginManifest,
) -> DevicePluginRuntimeContribution:
    entrypoint = manifest.spec.runtime.backendEntrypoint
    if manifest.spec.runtime.isolation != "transitional-in-process":
        raise DevicePluginCompositionError(
            "This Mission Core build only supports reviewed transitional-in-process "
            f"plugins; {manifest.metadata.id} requests {manifest.spec.runtime.isolation}"
        )
    factory = _load_factory(entrypoint)
    contribution = factory(repository_root)
    if not isinstance(contribution, DevicePluginRuntimeContribution):
        raise DevicePluginCompositionError(
            f"Device-plugin factory {entrypoint} returned an invalid contribution"
        )

    try:
        runtime = contribution.runtime
        descriptor = runtime.descriptor
        if descriptor.plugin_id != manifest.metadata.id:
            raise DevicePluginCompositionError(
                "Device-plugin manifest/runtime id mismatch: "
                f"{manifest.metadata.id} != {descriptor.plugin_id}"
            )
        if descriptor.plugin_version != manifest.metadata.version:
            raise DevicePluginCompositionError(
                "Device-plugin manifest/runtime version mismatch: "
                f"{manifest.metadata.version} != {descriptor.plugin_version}"
            )
        if manifest.apiVersion not in descriptor.supported_host_api_versions:
            raise DevicePluginCompositionError(
                f"Device-plugin runtime does not support {manifest.apiVersion}"
            )
        manifest_actions = frozenset(action.id for action in manifest.spec.actions)
        if frozenset(descriptor.action_ids) != manifest_actions:
            raise DevicePluginCompositionError(
                "Device-plugin manifest/runtime actions mismatch for "
                f"{descriptor.plugin_id}"
            )
        observation = contribution.observation
        if observation is not None:
            if not observation.archives:
                raise DevicePluginCompositionError(
                    "Device-plugin observation contribution is empty for "
                    f"{descriptor.plugin_id}"
                )
            archive_ids: set[str] = set()
            archive_roots: set[Path] = set()
            for archive in observation.archives:
                if archive.plugin_id != descriptor.plugin_id:
                    raise DevicePluginCompositionError(
                        "Device-plugin observation/runtime id mismatch: "
                        f"{archive.plugin_id} != {descriptor.plugin_id}"
                    )
                if archive.archive_id in archive_ids:
                    raise DevicePluginCompositionError(
                        f"Duplicate observation archive id for {descriptor.plugin_id}"
                    )
                resolved_root = archive.root.expanduser().resolve()
                if resolved_root in archive_roots:
                    raise DevicePluginCompositionError(
                        f"Duplicate observation archive root for {descriptor.plugin_id}"
                    )
                archive_ids.add(archive.archive_id)
                archive_roots.add(resolved_root)
        handshake_request = RuntimeHandshakeRequest(
            handshake_id=uuid4().hex,
            plugin_id=manifest.metadata.id,
            plugin_version=manifest.metadata.version,
            host_api_version=manifest.apiVersion,
            requested_at=datetime.now(UTC),
            required_action_ids=tuple(sorted(manifest_actions)),
        )
        try:
            handshake = runtime.handshake(handshake_request)
            runtime_health = runtime.health()
        except Exception as exc:
            raise DevicePluginCompositionError(
                f"Device-plugin runtime handshake failed for {descriptor.plugin_id}"
            ) from exc
        if (
            handshake.handshake_id != handshake_request.handshake_id
            or handshake.descriptor != descriptor
            or handshake.accepted_host_api_version != manifest.apiVersion
            or runtime_health.runtime_instance_id != handshake.runtime_instance_id
            or runtime_health.plugin_id != descriptor.plugin_id
            or runtime_health.plugin_version != descriptor.plugin_version
            or runtime_health.runtime_protocol_version
            != descriptor.runtime_protocol_version
            or runtime_health.status != "ready"
        ):
            raise DevicePluginCompositionError(
                f"Device-plugin runtime handshake is inconsistent for {descriptor.plugin_id}"
            )
    except Exception as exc:
        _add_cleanup_notes(exc, _close_contributions((contribution,)))
        raise
    return contribution


def load_installed_device_plugins(
    repository_root: Path,
) -> InstalledDevicePluginEnvironment:
    """Load local manifests and admit only compatible plugin runtimes."""

    catalog = DevicePluginCatalog(repository_root)
    manifests = catalog.manifests()
    loaded: list[DevicePluginRuntimeContribution] = []
    try:
        for manifest in manifests:
            loaded.append(_load_contribution(repository_root, manifest))
    except Exception as exc:
        _add_cleanup_notes(exc, _close_contributions(loaded))
        raise
    contributions = tuple(loaded)
    dispatcher = DevicePluginDispatcher(
        [contribution.runtime for contribution in contributions]
    )
    catalog_ids = {manifest.metadata.id for manifest in manifests}
    if set(dispatcher.action_declarations) != catalog_ids:
        raise DevicePluginCompositionError(
            "Device-plugin catalog/runtime composition is incomplete"
        )
    return InstalledDevicePluginEnvironment(
        catalog=catalog,
        dispatcher=dispatcher,
        legacy_routers=tuple(
            router for contribution in contributions for router in contribution.legacy_routers
        ),
        _contributions=contributions,
    )
