"""Canonical capability store (0.1.19) — UI ↔ canonical ↔ backend mapping layer."""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

from ..capability_manager import CAPABILITY_DEFINITIONS, CapabilityStateManager
from .action_session_store import ActionSessionStore
from .command_models import CANONICAL_CAPABILITY_KEYS

SCHEMA_VERSION = "0.1.21"

# Canonical capability metadata (12 caps).
CANONICAL_DEFINITIONS: dict[str, dict[str, Any]] = {
    "READ_ONLY_SYSTEM_INFO": {
        "risk_level": "low",
        "require_confirmation": False,
        "audit_required": True,
    },
    "READ_FILES_ALLOWED_PATHS": {
        "risk_level": "medium",
        "require_confirmation": True,
        "audit_required": True,
    },
    "WRITE_FILES_ALLOWED_PATHS": {
        "risk_level": "medium",
        "require_confirmation": True,
        "audit_required": True,
    },
    "RUN_SAFE_COMMANDS": {
        "risk_level": "high",
        "require_confirmation": True,
        "audit_required": True,
    },
    "INSTALL_PACKAGES": {
        "risk_level": "high",
        "require_confirmation": True,
        "audit_required": True,
    },
    "RESTART_SERVICES": {
        "risk_level": "high",
        "require_confirmation": True,
        "audit_required": True,
    },
    "MANAGE_DESKTOP_APPS": {
        "risk_level": "high",
        "require_confirmation": True,
        "audit_required": True,
    },
    "BROWSER_OPEN_URL": {
        "risk_level": "low",
        "require_confirmation": True,
        "audit_required": True,
    },
    "PROJECT_WORKSPACE_ACCESS": {
        "risk_level": "medium",
        "require_confirmation": True,
        "audit_required": True,
    },
    "AI_ASSISTED_EXECUTION": {
        "risk_level": "meta",
        "require_confirmation": True,
        "audit_required": True,
    },
    "LOCAL_MODEL_EXECUTION": {
        "risk_level": "medium",
        "require_confirmation": True,
        "audit_required": True,
    },
    "OPENAI_ASSISTED_EXECUTION": {
        "risk_level": "medium",
        "require_confirmation": True,
        "audit_required": True,
    },
    "SCREEN_CAPTURE": {
        "risk_level": "high",
        "require_confirmation": True,
        "audit_required": True,
    },
}

# UI key → canonical capabilities enabled when that UI row is granted.
UI_KEY_TO_CANONICAL: dict[str, frozenset[str]] = {
    "file_access": frozenset(
        {
            "READ_FILES_ALLOWED_PATHS",
            "WRITE_FILES_ALLOWED_PATHS",
            "PROJECT_WORKSPACE_ACCESS",
        }
    ),
    "application_control": frozenset({"MANAGE_DESKTOP_APPS", "RUN_SAFE_COMMANDS"}),
    "mouse_control": frozenset({"MANAGE_DESKTOP_APPS"}),
    "keyboard_control": frozenset({"MANAGE_DESKTOP_APPS"}),
    "network_control": frozenset({"READ_ONLY_SYSTEM_INFO", "RESTART_SERVICES", "BROWSER_OPEN_URL"}),
    "software_management": frozenset({"INSTALL_PACKAGES", "RESTART_SERVICES"}),
    "root_operations": frozenset({"RUN_SAFE_COMMANDS"}),
    "screen_vision": frozenset({"SCREEN_CAPTURE"}),
}

# Canonical → UI keys that can grant it (any enabled UI key is sufficient).
CANONICAL_TO_UI_KEYS: dict[str, frozenset[str]] = {
    "READ_ONLY_SYSTEM_INFO": frozenset({"network_control"}),
    "READ_FILES_ALLOWED_PATHS": frozenset({"file_access"}),
    "WRITE_FILES_ALLOWED_PATHS": frozenset({"file_access"}),
    "RUN_SAFE_COMMANDS": frozenset({"application_control", "root_operations"}),
    "INSTALL_PACKAGES": frozenset({"software_management"}),
    "RESTART_SERVICES": frozenset({"software_management", "network_control"}),
    "MANAGE_DESKTOP_APPS": frozenset(
        {"application_control", "mouse_control", "keyboard_control"}
    ),
    "BROWSER_OPEN_URL": frozenset({"network_control", "application_control"}),
    "PROJECT_WORKSPACE_ACCESS": frozenset({"file_access"}),
    "AI_ASSISTED_EXECUTION": frozenset(),
    "LOCAL_MODEL_EXECUTION": frozenset(),
    "OPENAI_ASSISTED_EXECUTION": frozenset(),
    "SCREEN_CAPTURE": frozenset({"screen_vision"}),
}


class CapabilityStore:
    """Canonical capability layer over ``CapabilityStateManager`` (8 UI keys).

    Implements ``CapabilityGrantChecker`` for ``PolicyEngine``.
    Authority source remains local UI grants persisted by the manager.
    """

    def __init__(self, manager: CapabilityStateManager | None = None) -> None:
        self._manager = manager or CapabilityStateManager()
        self._action_session_store = ActionSessionStore()

    def _refresh_manager_state(self) -> None:
        self._manager.reload_if_stale()

    @property
    def action_session_store(self) -> ActionSessionStore:
        return self._action_session_store

    @property
    def manager(self) -> CapabilityStateManager:
        return self._manager

    # --- CapabilityGrantChecker (PolicyEngine) --------------------------------

    def is_granted(self, capability_key: str, *, scope: dict[str, Any] | None = None) -> bool:
        self._refresh_manager_state()
        if capability_key not in CANONICAL_CAPABILITY_KEYS:
            return False
        if self.is_expired(capability_key):
            return False

        ui_keys = CANONICAL_TO_UI_KEYS.get(capability_key, frozenset())
        if not ui_keys:
            return False

        for ui_key in ui_keys:
            cap = self._manager.get(ui_key)
            if cap is None or not cap.enabled:
                continue
            if cap.expiry_type == "one_shot" and cap.one_shot_remaining is not None:
                if cap.one_shot_remaining <= 0:
                    continue
            if capability_key in UI_KEY_TO_CANONICAL.get(ui_key, frozenset()):
                return True
        return False

    def is_expired(self, capability_key: str) -> bool:
        self._refresh_manager_state()
        ui_keys = CANONICAL_TO_UI_KEYS.get(capability_key, frozenset())
        if not ui_keys:
            return True

        any_granted = False
        all_expired = True
        for ui_key in ui_keys:
            cap = self._manager.get(ui_key)
            if cap is None or not cap.enabled:
                continue
            any_granted = True
            if cap.is_expired():
                continue
            if cap.expiry_type == "one_shot" and cap.one_shot_remaining is not None:
                if cap.one_shot_remaining <= 0:
                    continue
            all_expired = False

        if not any_granted:
            return False
        return all_expired

    # --- Lifecycle ------------------------------------------------------------

    def consume_one_shot(self, canonical_cap: str) -> None:
        """Disable one-shot UI grants backing *canonical_cap* after successful exec."""
        ui_keys = CANONICAL_TO_UI_KEYS.get(canonical_cap, frozenset())
        for ui_key in ui_keys:
            cap = self._manager.get(ui_key)
            if cap is None or not cap.enabled or cap.expiry_type != "one_shot":
                continue
            if canonical_cap not in UI_KEY_TO_CANONICAL.get(ui_key, frozenset()):
                continue
            self._manager.consume_one_shot(ui_key)

    def revoke_all(self) -> None:
        """Revoke every UI capability and record ``revoked_all_at``."""
        for ui_key in CAPABILITY_DEFINITIONS:
            self._manager.revoke(ui_key)
        self._manager.set_revoked_all_at(datetime.now(timezone.utc).isoformat())
        self._action_session_store.revoke_all()

    @property
    def always_allow_types(self) -> frozenset[str]:
        return frozenset(self._manager.always_allow_types)

    def add_always_allow_type(self, command_type: str) -> bool:
        return self._manager.add_always_allow_type(command_type)

    def read_canonical_counts(self) -> tuple[int, int]:
        snap = self.snapshot_for_backend()
        granted = int(snap.get("granted_canonical_count") or 0)
        total = int(snap.get("total_canonical_count") or len(CANONICAL_CAPABILITY_KEYS))
        return granted, total

    # --- UI delegation --------------------------------------------------------

    def grant_ui(self, ui_key: str, *, expiry_preset: str = "permanent"):
        return self._manager.grant(ui_key, expiry_preset=expiry_preset)

    def revoke_ui(self, ui_key: str):
        return self._manager.revoke(ui_key)

    def read_ui_snapshot(self) -> dict[str, dict[str, Any]]:
        self._refresh_manager_state()
        return {key: cap.to_dict() for key, cap in self._manager.all().items()}

    def is_ui_enabled(self, ui_key: str) -> bool:
        return self._manager.is_enabled(ui_key)

    # --- Backend evidence -----------------------------------------------------

    def snapshot_for_backend(self) -> dict[str, Any]:
        """Evidence payload: UI state + canonical grant view."""
        base = self._manager.snapshot_for_backend()
        canonical_grants: dict[str, dict[str, Any]] = {}
        for key in CANONICAL_CAPABILITY_KEYS:
            canonical_grants[key] = {
                "granted": self.is_granted(key),
                "expired": self.is_expired(key),
                "risk_level": CANONICAL_DEFINITIONS.get(key, {}).get("risk_level", "medium"),
            }
        base["schema_version"] = SCHEMA_VERSION
        base["canonical_capabilities"] = canonical_grants
        base["granted_canonical_count"] = sum(1 for v in canonical_grants.values() if v["granted"])
        base["total_canonical_count"] = len(CANONICAL_CAPABILITY_KEYS)
        return base

    def canonical_keys_for_ui(self, ui_key: str) -> frozenset[str]:
        return UI_KEY_TO_CANONICAL.get(ui_key, frozenset())

    def ui_keys_for_canonical(self, canonical_key: str) -> frozenset[str]:
        return CANONICAL_TO_UI_KEYS.get(canonical_key, frozenset())
