# platform/deskclaw/setup/loader.py
# ═══════════════════════════════════════════════════════════════
# Phase 2: Setup Loader — Plugin Configuration & Session Init
# Parses deskclaw-plugin.yaml, resolves environment variables,
# validates Agent Adapters, and emits session_config.json.
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

import hashlib
import json
import os
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

try:
    import yaml  # type: ignore[import-untyped]
except ModuleNotFoundError:  # pragma: no cover — yaml is optional in minimal installs
    yaml = None  # type: ignore[assignment]

from deskclaw.adapters.registry import get_adapter, list_adapters

# ── Constants ──

_ENV_VAR_PATTERN = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")


# ── Exceptions ──


class SetupError(Exception):
    """Raised when Phase 2 setup fails."""


class EnvVarMissingError(SetupError):
    """A referenced environment variable is not set."""


class AdapterValidationError(SetupError):
    """The requested adapter cannot be instantiated."""


# ── Data Models ──


@dataclass
class PluginConfig:
    """Parsed representation of ``deskclaw-plugin.yaml``."""

    metadata: dict[str, str] = field(default_factory=dict)
    agent: dict[str, Any] = field(default_factory=dict)
    backup: dict[str, Any] = field(default_factory=dict)
    execution: dict[str, Any] = field(default_factory=dict)
    evaluation: dict[str, Any] = field(default_factory=dict)
    output: dict[str, Any] = field(default_factory=dict)

    @property
    def default_agent(self) -> str:
        return self.agent.get("default", "self")

    @property
    def agent_configs(self) -> dict[str, dict]:
        return self.agent.get("config", {})

    @property
    def backup_paths(self) -> list[str]:
        return self.backup.get("paths", [])

    @property
    def backup_exclude(self) -> list[str]:
        return self.backup.get(
            "exclude", [".git", "__pycache__", "node_modules", "*.pyc"]
        )

    @property
    def snapshot_dir(self) -> str:
        return self.backup.get("snapshot_dir", "./snapshots")

    @property
    def max_snapshots(self) -> int:
        return self.backup.get("max_snapshots", 20)

    @property
    def benchmarks(self) -> list[dict]:
        return self.execution.get("benchmarks", [])

    @property
    def parallel_config(self) -> dict[str, Any]:
        return self.execution.get("parallel", {})

    @property
    def results_dir(self) -> str:
        return self.output.get("results_dir", "./results")


@dataclass
class SessionConfig:
    """Output of Phase 2: everything needed to start Phase 3."""

    snapshot_id: str
    agent: str
    plugin_yaml_sha256: str
    mock_enabled: bool
    benchmarks_loaded: list[str]
    parallel_workers: int
    created_at: str  # ISO-8601

    def to_dict(self) -> dict:
        return {
            "snapshot_id": self.snapshot_id,
            "agent": self.agent,
            "plugin_yaml_sha256": self.plugin_yaml_sha256,
            "mock_enabled": self.mock_enabled,
            "benchmarks_loaded": self.benchmarks_loaded,
            "parallel_workers": self.parallel_workers,
            "created_at": self.created_at,
        }


# ── Core Loader ──


class SetupLoader:
    """
    Phase 2 Setup engine.

    1. Parse ``deskclaw-plugin.yaml``
    2. Resolve ``${ENV_VAR}`` references
    3. Validate Agent Adapter registration & instantiation
    4. Emit ``session_config.json`` to ``results/{snapshot_id}/``
    """

    def __init__(
        self, plugin_path: str | Path = "plugins/deskclaw-plugin.yaml"
    ) -> None:
        self.plugin_path = Path(plugin_path)

    # ── Public API ──

    def load(
        self,
        *,
        agent_override: str | None = None,
        snapshot_id: str | None = None,
        results_base: str | Path | None = None,
    ) -> SessionConfig:
        """
        Full Phase 2 setup flow.

        Args:
            agent_override: Use this adapter instead of the YAML default.
            snapshot_id:    Bind to an existing Phase 1 snapshot ID.
            results_base:   Root directory for results output.

        Returns:
            :class:`SessionConfig` written to disk and returned.

        Raises:
            SetupError: On YAML parse failure, missing env vars, or adapter issues.
        """
        raw = self._read_yaml()
        resolved = self._resolve_env_vars_selective(raw, agent_override=agent_override)
        config = self._build_plugin_config(resolved)

        agent_name = agent_override or config.default_agent
        self._validate_adapter(agent_name, config)

        yaml_sha = _sha256_string(self.plugin_path.read_text())

        now = datetime.now(timezone.utc)
        sid = snapshot_id or now.strftime("%Y%m%dT%H%M%S") + f"_{agent_name}"

        benchmarks_loaded = [b.get("name", "unknown") for b in config.benchmarks]
        parallel_workers = config.parallel_config.get("max_workers", 4)
        mock_enabled = config.execution.get("mock", {}).get("enabled", False)

        session = SessionConfig(
            snapshot_id=sid,
            agent=agent_name,
            plugin_yaml_sha256=yaml_sha,
            mock_enabled=mock_enabled,
            benchmarks_loaded=benchmarks_loaded,
            parallel_workers=parallel_workers,
            created_at=now.isoformat(),
        )

        results_root = Path(results_base) if results_base else Path(config.results_dir)
        self._write_session_config(session, results_root)
        return session

    def parse_only(self, *, agent_hint: str | None = None) -> PluginConfig:
        """Parse and resolve the YAML without validating adapters or writing files."""
        raw = self._read_yaml()
        if agent_hint is not None:
            resolved = self._resolve_env_vars_selective(raw, agent_override=agent_hint)
        else:
            resolved = self._resolve_env_vars(raw)
        return self._build_plugin_config(resolved)

    # ── YAML Parsing ──

    def _read_yaml(self) -> dict:
        """Read and parse the plugin YAML file."""
        if not self.plugin_path.exists():
            raise SetupError(f"Plugin config not found: {self.plugin_path}")

        text = self.plugin_path.read_text(encoding="utf-8")

        if yaml is not None:
            try:
                data = yaml.safe_load(text)
            except yaml.YAMLError as exc:
                raise SetupError(f"YAML parse error: {exc}") from exc
        else:
            # Minimal fallback — try JSON (e.g. for test fixtures)
            try:
                data = json.loads(text)
            except json.JSONDecodeError as exc:
                raise SetupError(
                    "PyYAML not installed and file is not valid JSON. "
                    "Install PyYAML: pip install pyyaml"
                ) from exc

        if not isinstance(data, dict):
            raise SetupError(f"Expected top-level dict, got {type(data).__name__}")
        return data

    # ── Environment Variable Resolution ──

    def _resolve_env_vars_selective(
        self, data: Any, *, agent_override: str | None = None
    ) -> Any:
        """
        Resolve ``${VAR_NAME}`` tokens selectively.

        For ``agent.config``, only the selected agent's config block is resolved.
        Other adapters' configs are left with raw ``${…}`` tokens intact so that
        missing env vars for unused adapters do not cause failures.

        Args:
            data:           Raw parsed config dict.
            agent_override: If provided, only this agent's config is env-resolved.
                            If ``None``, the ``agent.default`` value is used.
        """
        if not isinstance(data, dict):
            return self._resolve_env_vars(data)

        result: dict = {}
        for key, value in data.items():
            if key == "agent" and isinstance(value, dict):
                # Determine which agent to resolve env vars for
                selected = agent_override or value.get("default", "self")
                agent_block = dict(value)
                configs = agent_block.get("config", {})
                if isinstance(configs, dict):
                    resolved_configs: dict = {}
                    for adapter_name, adapter_cfg in configs.items():
                        if adapter_name == selected:
                            resolved_configs[adapter_name] = self._resolve_env_vars(
                                adapter_cfg
                            )
                        else:
                            # Keep raw — don't resolve env vars for unused adapters
                            resolved_configs[adapter_name] = adapter_cfg
                    agent_block["config"] = resolved_configs
                # Resolve non-config parts of the agent block (e.g. "default")
                for ak, av in agent_block.items():
                    if ak != "config":
                        agent_block[ak] = self._resolve_env_vars(av)
                result[key] = agent_block
            else:
                result[key] = self._resolve_env_vars(value)
        return result

    def _resolve_env_vars(self, data: Any) -> Any:
        """
        Recursively replace ``${VAR_NAME}`` tokens with ``os.environ[VAR_NAME]``.

        Raises:
            EnvVarMissingError: If a referenced variable is not set.
        """
        if isinstance(data, str):
            return self._resolve_string(data)
        if isinstance(data, dict):
            return {k: self._resolve_env_vars(v) for k, v in data.items()}
        if isinstance(data, list):
            return [self._resolve_env_vars(item) for item in data]
        return data

    @staticmethod
    def _resolve_string(value: str) -> str:
        """Replace all ``${VAR}`` occurrences in a single string."""

        def _replacer(match: re.Match) -> str:
            var = match.group(1)
            env_val = os.environ.get(var)
            if env_val is None:
                raise EnvVarMissingError(
                    f"Environment variable '${{{var}}}' is not set. "
                    f"Set it or add it to your .env file."
                )
            return env_val

        return _ENV_VAR_PATTERN.sub(_replacer, value)

    # ── Config Construction ──

    @staticmethod
    def _build_plugin_config(data: dict) -> PluginConfig:
        return PluginConfig(
            metadata=data.get("metadata", {}),
            agent=data.get("agent", {}),
            backup=data.get("backup", {}),
            execution=data.get("execution", {}),
            evaluation=data.get("evaluation", {}),
            output=data.get("output", {}),
        )

    # ── Adapter Validation ──

    @staticmethod
    def _validate_adapter(agent_name: str, config: PluginConfig) -> None:
        """
        Ensure the adapter is registered and can be instantiated.

        Raises:
            AdapterValidationError: On failure.
        """
        available = list_adapters()
        if agent_name not in available:
            raise AdapterValidationError(
                f"Adapter '{agent_name}' is not registered. Available: {available}"
            )

        agent_cfg = config.agent_configs.get(agent_name, {})
        try:
            adapter = get_adapter(agent_name, **agent_cfg)
        except Exception as exc:
            raise AdapterValidationError(
                f"Failed to instantiate adapter '{agent_name}': {exc}"
            ) from exc

        # Quick sanity — adapter must respond to name()
        try:
            _ = adapter.name()
        except Exception as exc:
            raise AdapterValidationError(
                f"Adapter '{agent_name}' instantiated but name() failed: {exc}"
            ) from exc

    # ── Output ──

    @staticmethod
    def _write_session_config(session: SessionConfig, results_root: Path) -> Path:
        """Write ``session_config.json`` into ``results_root/{snapshot_id}/``."""
        output_dir = results_root / session.snapshot_id
        output_dir.mkdir(parents=True, exist_ok=True)
        out_path = output_dir / "session_config.json"
        out_path.write_text(json.dumps(session.to_dict(), indent=2))
        return out_path


# ── Utilities ──


def _sha256_string(text: str) -> str:
    """SHA256 hex digest of a UTF-8 string."""
    return hashlib.sha256(text.encode("utf-8")).hexdigest()
