"""Manifest client - lifecycle management for training manifests.

Handles creating, loading, saving, and submitting manifests to the training queue.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal

from appinfra import DotDict
from appinfra.log import Logger

from ..schema import Adapter, SubmitResult
from ..storage import Storage
from .loader import load_manifest as _load_manifest
from .loader import save_manifest as _save_manifest
from .schema import Data, Deployment, Manifest, Source

if TYPE_CHECKING:
    pass

# Keys that belong in training config (vs method_config or lora)
_TRAINING_KEYS = frozenset(
    [
        "num_epochs",
        "batch_size",
        "gradient_accumulation_steps",
        "learning_rate",
        "warmup_ratio",
        "max_seq_length",
        "max_grad_norm",
        "logging_steps",
        "save_steps",
        "eval_split",
        "fp16",
        "bf16",
        "gradient_checkpointing",
        "seed",
    ]
)

# Map profile keys to training config keys
_PROFILE_KEY_MAP = {
    "epochs": "num_epochs",
}


class Client:
    """Client for training manifest lifecycle management.

    Provides methods to create, load, save, and submit manifests.
    Accessed via Client.train.manifest.

    Usage:
        # Create a manifest
        manifest = kelt.train.manifest.create(
            adapter="coding-v1",
            method="dpo",
            model="Qwen/Qwen2.5-7B-Instruct",
            data=[{"prompt": "...", "chosen": "...", "rejected": "..."}],
        )

        # Save to file
        kelt.train.manifest.save(manifest, Path("coding-v1.yaml"))

        # Submit to training queue
        kelt.train.manifest.submit(manifest)

        # List pending manifests
        kelt.train.manifest.list_pending()
    """

    def __init__(
        self,
        lg: Logger,
        storage: Storage,
        default_profiles: dict[str, dict] | None = None,
    ) -> None:
        """Initialize manifest client.

        Args:
            lg: Logger instance.
            storage: Storage instance for filesystem operations.
            default_profiles: Default training profiles by method.
        """
        self._lg = lg
        self._storage = storage
        self._default_profiles = default_profiles or {}
        self._manifest_cache: dict[str, Manifest] = {}  # md5 -> Manifest cache

    def _extract_nested_configs(
        self, config: dict[str, Any]
    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
        """Extract nested lora/training sections from config, returning flat remainder."""
        config = dict(config) if config else {}

        # Validate lora/training are dicts if present
        if "lora" in config and not isinstance(config["lora"], dict):
            raise ValueError(f"config.lora must be a dict, got {type(config['lora']).__name__}")
        if "training" in config and not isinstance(config["training"], dict):
            raise ValueError(
                f"config.training must be a dict, got {type(config['training']).__name__}"
            )

        nested_lora = config.pop("lora", {})
        nested_training = config.pop("training", {})
        return config, nested_lora, nested_training

    def _build_manifest_configs(
        self,
        method: Literal["dpo", "sft", "prompt"],
        model: str | None,
        config: dict[str, Any] | None,
    ) -> tuple[DotDict, DotDict, DotDict]:
        """Build configuration objects by merging default profile with agent overrides."""
        defaults = dict(self._default_profiles.get(method, {}))
        flat_config, nested_lora, nested_training = self._extract_nested_configs(config or {})

        # Merge flat config with defaults
        merged = {**defaults, **flat_config}

        # Map profile keys to training config keys (e.g., epochs -> num_epochs)
        for profile_key, config_key in _PROFILE_KEY_MAP.items():
            if profile_key in merged and config_key not in merged:
                merged[config_key] = merged.pop(profile_key)

        # Build training config: defaults + flat keys + nested training overrides
        training_config = DotDict({k: merged[k] for k in _TRAINING_KEYS if k in merged})
        training_config.update(nested_training)
        if model:
            training_config["requested_model"] = model

        # Build lora config: defaults + nested lora overrides
        lora_config = DotDict(merged.get("lora", {}))
        lora_config.update(nested_lora)

        method_config = self._build_method_config(method, merged)
        return lora_config, training_config, method_config

    def _build_method_config(
        self, method: Literal["dpo", "sft", "prompt"], merged: dict[str, Any]
    ) -> DotDict:
        """Build method-specific config (beta for DPO, virtual tokens for prompt, etc.)."""
        if method == "dpo":
            return DotDict(self._extract_method_config(merged, ["beta"]))
        if method == "prompt":
            keys = ["num_virtual_tokens", "prompt_tuning_init", "prompt_tuning_init_text"]
            return DotDict(self._extract_method_config(merged, keys))
        return DotDict()

    def create(
        self,
        adapter: str,
        method: Literal["dpo", "sft", "prompt"],
        data: list[dict[str, Any]],
        *,
        model: str | None = None,
        parent: Adapter | None = None,
        context_key: str | None = None,
        schema_name: str | None = None,
        description: str | None = None,
        config: dict[str, Any] | None = None,
        deployment_policy: Literal["skip", "add", "replace"] | None = None,
    ) -> Manifest:
        """Create a new training manifest with inline data.

        Args:
            adapter: Output adapter key (series name, e.g., "my-agent-sft").
            method: Training method ("dpo", "sft", or "prompt").
            data: List of training records (inline data).
            model: Optional base model (can be specified at training time via --model).
            parent: Parent adapter for lineage (continue training from this).
            context_key: Agent context key for provenance.
            schema_name: Database schema where training data originated.
            description: Human-readable description.
            config: Override training configuration (num_epochs, etc.).
            deployment_policy: Deployment policy after training:
                - "skip": Don't deploy the adapter.
                - "add": Deploy as new version, keep existing versions.
                - "replace": Deploy and remove existing versions (default).

        Returns:
            Manifest instance.
        """
        source = Source(context_key=context_key, schema_name=schema_name, description=description)
        lora_config, training_config, method_config = self._build_manifest_configs(
            method, model, config
        )

        # Build deployment config (default to "replace" if not specified)
        deployment = Deployment(policy=deployment_policy or "replace")

        manifest = Manifest(
            adapter=adapter,
            method=method,
            data=Data(format="inline", records=data),
            deployment=deployment,
            source=source,
            parent=parent,
            lora=lora_config,
            training=training_config,
            method_config=method_config,
        )

        self._lg.info("created manifest", extra={"adapter": adapter, "method": method})
        return manifest

    def _extract_method_config(self, overrides: dict[str, Any], keys: list[str]) -> dict[str, Any]:
        """Extract method-specific config values from overrides."""
        return {k: overrides[k] for k in keys if k in overrides}

    def load(self, path: Path) -> Manifest:
        """Load a manifest from file.

        Args:
            path: Path to manifest YAML file.

        Returns:
            Manifest instance.
        """
        return _load_manifest(path)

    def save(self, manifest: Manifest, path: Path) -> None:
        """Save a manifest to file.

        Args:
            manifest: Manifest to save.
            path: Output path for YAML file.
        """
        _save_manifest(manifest, path)
        self._lg.info("saved manifest", extra={"path": str(path)})

    def submit(self, manifest: Manifest) -> SubmitResult:
        """Submit a manifest to the pending training queue.

        Saves the manifest to the registry's pending directory.

        Args:
            manifest: Manifest to submit.

        Returns:
            SubmitResult with adapter key, timestamp, and storage location.

        Raises:
            ValueError: If manifest with same key already pending.
        """
        result = self._storage.submit_manifest(manifest)
        self._lg.info("submitted manifest", extra={"adapter": manifest.adapter})
        return result

    def list_pending(self) -> list[Manifest]:
        """List manifests waiting for training.

        Returns:
            List of pending manifests.
        """
        return self._storage.list_pending_manifests()

    def list_completed(self) -> list[Manifest]:
        """List completed manifests.

        Returns:
            List of completed manifests.
        """
        return self._storage.list_completed_manifests()

    def _validate_adapter_name(self, adapter: str) -> None:
        """Validate adapter name has no path traversal characters."""
        self._storage.validate_key(adapter)

    def get_pending(self, adapter: str) -> Manifest | None:
        """Get a pending manifest by adapter key.

        Args:
            adapter: Adapter key to look up.

        Returns:
            Manifest or None if not found.
        """
        return self._storage.get_pending_manifest(adapter)

    def remove_pending(self, adapter: str) -> None:
        """Remove a manifest from the pending queue.

        Args:
            adapter: Adapter key of manifest to remove.

        Raises:
            FileNotFoundError: If manifest not in queue.
            ValueError: If adapter name contains path traversal characters.
        """
        self._storage.remove_pending_manifest(adapter)
        self._lg.info("removed from queue", extra={"adapter": adapter})

    def find_adapter(self, md5: str) -> Adapter | None:
        """Find an adapter by its md5 hash.

        Searches completed manifests for an adapter with matching md5.

        Args:
            md5: MD5 hash of adapter weights (12 char hex).

        Returns:
            Adapter if found, None otherwise.
        """
        manifest = self._storage.find_adapter_by_md5(md5)
        if manifest and manifest.output and manifest.output.adapter:
            adapter: Adapter = manifest.output.adapter
            return adapter
        return None

    def get_manifest(self, md5: str) -> Manifest | None:
        """Get a completed manifest by adapter md5.

        Use this to look up the manifest for a deployed adapter, e.g., to
        determine which schema the adapter's training data came from.

        Results are cached since manifests don't change at runtime.

        Args:
            md5: MD5 hash of adapter weights (32 char hex).

        Returns:
            Manifest if found, None otherwise.

        Example:
            manifest = client.get_manifest("e6a8a798834d1c5f9b2a3e4d5f6a7b8c")
            schema = manifest.source.schema_name if manifest and manifest.source else None
        """
        if md5 in self._manifest_cache:
            return self._manifest_cache[md5]

        manifest = self._storage.find_adapter_by_md5(md5)
        if manifest:
            self._manifest_cache[md5] = manifest
        return manifest

    def get_latest_completed(
        self, adapter: str | None = None, context_key: str | None = None
    ) -> Manifest | None:
        """Get the most recently completed manifest.

        Args:
            adapter: Filter by adapter key.
            context_key: Filter by agent context key.

        Returns:
            Most recent completed manifest matching filters, or None.
        """
        manifests: list[Manifest] = []
        for manifest in self.list_completed():
            if adapter and manifest.adapter != adapter:
                continue
            if context_key and (not manifest.source or manifest.source.context_key != context_key):
                continue
            if manifest.output and manifest.output.status == "completed":
                manifests.append(manifest)

        if not manifests:
            return None

        # Sort by completed_at descending, return most recent
        return max(
            manifests,
            key=lambda m: (
                m.output.completed_at if m.output and m.output.completed_at else m.created_at
            ),
        )
