"""Deterministic local selector for adaptive expansion recipes."""

from __future__ import annotations

import hashlib
import json
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional

from model_osmosis.optimizer.self_optimize import (
    record_subsystem_selection,
    resolve_subsystem_mode,
)


EXPAND_ENGINE_MODES = {"auto", "off"}
EXPAND_SELECTOR_VERSION = "expand-selector-v1"
EXPAND_PACK_VERSION = "builtin-expand-pack-v1"
EXPAND_SUBSYSTEM = "expand"

_RECIPES: Dict[str, Dict[str, Any]] = {
    "expand_balanced_v1": {
        "width_factor": 1.25,
        "depth_layers": 1,
        "adaptation_epochs": 2,
        "task_bias": {
            "classification": 0.60,
            "regression": 0.60,
            "text": 0.55,
            "vision": 0.58,
            "audio": 0.56,
            "code": 0.56,
            "auto": 0.55,
        },
    },
    "expand_wider_v1": {
        "width_factor": 1.50,
        "depth_layers": 0,
        "adaptation_epochs": 1,
        "task_bias": {
            "classification": 0.64,
            "regression": 0.62,
            "text": 0.52,
            "vision": 0.63,
            "audio": 0.59,
            "code": 0.53,
            "auto": 0.54,
        },
    },
    "expand_deeper_v1": {
        "width_factor": 1.10,
        "depth_layers": 2,
        "adaptation_epochs": 3,
        "task_bias": {
            "classification": 0.57,
            "regression": 0.59,
            "text": 0.65,
            "vision": 0.55,
            "audio": 0.58,
            "code": 0.64,
            "auto": 0.56,
        },
    },
}


def _stable_float(seed: str) -> float:
    digest = hashlib.sha256(seed.encode("utf-8")).digest()
    raw = int.from_bytes(digest[:8], "big", signed=False)
    return raw / float(2**64 - 1)


def _context_fingerprint(context: Dict[str, Any]) -> str:
    normalized = json.dumps(context, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
    return hashlib.sha256(normalized.encode("utf-8")).hexdigest()


@dataclass(frozen=True)
class ExpandSelection:
    """Chosen expansion profile and deterministic selector metadata."""

    optimizer_mode: str
    profile_id: str
    recipe: Dict[str, Any]
    pack_version: str
    selector_version: str
    canary_status: str
    rollback_reason: str
    selection_latency_ms: float
    fallback_reason: str
    score: float

    def to_dict(self) -> Dict[str, Any]:
        return {
            "optimizer_mode": self.optimizer_mode,
            "profile_id": self.profile_id,
            "pack_version": self.pack_version,
            "selector_version": self.selector_version,
            "canary_status": self.canary_status,
            "rollback_reason": self.rollback_reason,
            "selection_latency_ms": round(float(self.selection_latency_ms), 4),
            "fallback_reason": self.fallback_reason,
            "score": round(float(self.score), 6),
            "recipe": dict(self.recipe),
        }


def resolve_expand_engine_mode(requested: Optional[str]) -> str:
    mode = resolve_subsystem_mode(requested, subsystem_env_var="OSMOSIS_EXPAND_OPTIMIZER")
    if mode not in EXPAND_ENGINE_MODES:
        raise ValueError(
            f"Unsupported expand optimizer mode '{mode}'. "
            f"Expected one of: {sorted(EXPAND_ENGINE_MODES)}"
        )
    return mode


def select_expand_profile(
    *,
    context: Dict[str, Any],
    optimizer_mode: Optional[str],
    seed: int,
    record_selection: bool = True,
) -> ExpandSelection:
    """Select an expansion profile with deterministic local ranking."""
    started = time.perf_counter()
    mode = resolve_expand_engine_mode(optimizer_mode)

    task = str(context.get("task", "auto") or "auto").strip().lower()
    fingerprint = _context_fingerprint(context)
    selected_profile = "expand_balanced_v1"
    selected_score = 0.0
    reason = ""

    if mode == "off":
        selected_profile = "expand_balanced_v1"
        selected_score = _stable_float(f"{fingerprint}:{seed}:expand_balanced_v1")
        reason = reason or "optimizer_off"
    else:
        ranked = []
        for profile_id, recipe in sorted(_RECIPES.items(), key=lambda item: item[0]):
            task_bias = float(
                (recipe.get("task_bias", {}) or {}).get(
                    task,
                    (recipe.get("task_bias", {}) or {}).get("auto", 0.5),
                )
            )
            jitter = _stable_float(f"{fingerprint}:{seed}:{profile_id}")
            score = (0.85 * task_bias) + (0.15 * jitter)
            ranked.append((score, profile_id))
        ranked.sort(key=lambda item: (-item[0], item[1]))
        selected_score, selected_profile = ranked[0]

    latency_ms = max(0.0, (time.perf_counter() - started) * 1000.0)
    metadata: Dict[str, Any]
    if record_selection:
        metadata = record_subsystem_selection(
            subsystem=EXPAND_SUBSYSTEM,
            profile_id=selected_profile,
            pack_version=EXPAND_PACK_VERSION,
            selector_version=EXPAND_SELECTOR_VERSION,
            selection_latency_ms=latency_ms,
            quality_ok=True,
            fallback_reason=reason,
            quality_score=selected_score,
        )
    else:
        metadata = {
            "pack_version": EXPAND_PACK_VERSION,
            "selector_version": EXPAND_SELECTOR_VERSION,
            "canary_status": "shadow",
            "rollback_reason": "",
            "selection_latency_ms": latency_ms,
        }
    recipe = dict(_RECIPES.get(selected_profile, _RECIPES["expand_balanced_v1"]))
    return ExpandSelection(
        optimizer_mode=mode,
        profile_id=selected_profile,
        recipe=recipe,
        pack_version=str(metadata.get("pack_version", EXPAND_PACK_VERSION) or EXPAND_PACK_VERSION),
        selector_version=str(metadata.get("selector_version", EXPAND_SELECTOR_VERSION) or EXPAND_SELECTOR_VERSION),
        canary_status=str(metadata.get("canary_status", "shadow")),
        rollback_reason=str(metadata.get("rollback_reason", "")),
        selection_latency_ms=float(metadata.get("selection_latency_ms", latency_ms) or latency_ms),
        fallback_reason=reason,
        score=float(selected_score),
    )


def finalize_expand_profile(
    selection: ExpandSelection,
    *,
    quality_ok: bool,
    fallback_reason: str = "",
    quality_score: Optional[float] = None,
) -> ExpandSelection:
    """Record post-run quality outcome and return updated selector metadata."""
    reason = str(fallback_reason or selection.fallback_reason or "")
    metadata = record_subsystem_selection(
        subsystem=EXPAND_SUBSYSTEM,
        profile_id=selection.profile_id,
        pack_version=selection.pack_version,
        selector_version=selection.selector_version,
        selection_latency_ms=float(selection.selection_latency_ms),
        quality_ok=bool(quality_ok),
        fallback_reason=reason,
        quality_score=(selection.score if quality_score is None else float(quality_score)),
    )
    return ExpandSelection(
        optimizer_mode=selection.optimizer_mode,
        profile_id=selection.profile_id,
        recipe=dict(selection.recipe),
        pack_version=str(metadata.get("pack_version", selection.pack_version) or selection.pack_version),
        selector_version=str(metadata.get("selector_version", selection.selector_version) or selection.selector_version),
        canary_status=str(metadata.get("canary_status", selection.canary_status)),
        rollback_reason=str(metadata.get("rollback_reason", selection.rollback_reason)),
        selection_latency_ms=float(metadata.get("selection_latency_ms", selection.selection_latency_ms) or selection.selection_latency_ms),
        fallback_reason=reason,
        score=float(selection.score if quality_score is None else quality_score),
    )
