"""Unified deterministic self-optimization state for local-only learning."""

from __future__ import annotations

import json
import os
import shutil
import time
from pathlib import Path
from typing import Any, Dict, Optional


OPTIMIZER_MODES = {"auto", "off"}
STATE_SCHEMA_VERSION = 1
DEFAULT_SUBSYSTEMS = (
    "strategy",
    "compression",
    "protection",
    "cross_merge_classic",
    "cross_merge_llm",
    "expand",
)
CANARY_STATUSES = {"shadow", "canary", "promoted", "rollback"}


def _utc_now() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def optimizer_home() -> Path:
    env = str(os.environ.get("OSMOSIS_OPTIMIZER_HOME", "")).strip()
    if env:
        return Path(env)
    return Path(".model_osmosis") / "self_optimize"


def state_file_path() -> Path:
    return optimizer_home() / "state.json"


def _subsystem_default() -> Dict[str, Any]:
    return {
        "canary_status": "shadow",
        "promoted_profile_id": "",
        "active_profile_id": "",
        "rollback_reason": "",
        "pack_version": "",
        "selector_version": "",
        "last_selection_latency_ms": 0.0,
        "runs": 0,
        "shadow_runs": 0,
        "canary_attempts": 0,
        "canary_passes": 0,
        "candidates": {},
        "scores": {},
    }


def _default_state() -> Dict[str, Any]:
    return {
        "schema_version": STATE_SCHEMA_VERSION,
        "updated_at": "",
        "rollback_marker": {},
        "subsystems": {name: _subsystem_default() for name in DEFAULT_SUBSYSTEMS},
    }


def _normalize_state(state: Dict[str, Any]) -> Dict[str, Any]:
    normalized = _default_state()
    normalized["schema_version"] = int(state.get("schema_version", STATE_SCHEMA_VERSION) or STATE_SCHEMA_VERSION)
    normalized["updated_at"] = str(state.get("updated_at", ""))
    rollback = state.get("rollback_marker", {})
    if isinstance(rollback, dict):
        normalized["rollback_marker"] = {
            "subsystem": str(rollback.get("subsystem", "")),
            "profile_id": str(rollback.get("profile_id", "")),
            "reason": str(rollback.get("reason", "")),
            "at": str(rollback.get("at", "")),
        }

    subsystems_in = state.get("subsystems", {})
    if isinstance(subsystems_in, dict):
        for name, value in subsystems_in.items():
            if not isinstance(name, str) or not isinstance(value, dict):
                continue
            entry = _subsystem_default()
            entry["canary_status"] = (
                str(value.get("canary_status", "shadow"))
                if str(value.get("canary_status", "shadow")) in CANARY_STATUSES
                else "shadow"
            )
            entry["promoted_profile_id"] = str(value.get("promoted_profile_id", ""))
            entry["active_profile_id"] = str(value.get("active_profile_id", ""))
            entry["rollback_reason"] = str(value.get("rollback_reason", ""))
            entry["pack_version"] = str(value.get("pack_version", ""))
            entry["selector_version"] = str(value.get("selector_version", ""))
            entry["last_selection_latency_ms"] = float(value.get("last_selection_latency_ms", 0.0) or 0.0)
            entry["runs"] = int(value.get("runs", 0) or 0)
            entry["shadow_runs"] = int(value.get("shadow_runs", 0) or 0)
            entry["canary_attempts"] = int(value.get("canary_attempts", 0) or 0)
            entry["canary_passes"] = int(value.get("canary_passes", 0) or 0)
            candidates = value.get("candidates", {})
            if isinstance(candidates, dict):
                entry["candidates"] = {
                    str(k): int(v or 0)
                    for k, v in sorted(candidates.items(), key=lambda item: str(item[0]))
                }
            scores = value.get("scores", {})
            if isinstance(scores, dict):
                entry["scores"] = {
                    str(k): float(v or 0.0)
                    for k, v in sorted(scores.items(), key=lambda item: str(item[0]))
                }
            normalized["subsystems"][str(name)] = entry
    return normalized


def load_state() -> Dict[str, Any]:
    path = state_file_path()
    if not path.exists():
        return _default_state()
    try:
        parsed = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return _default_state()
    if not isinstance(parsed, dict):
        return _default_state()
    return _normalize_state(parsed)


def save_state(state: Dict[str, Any]) -> None:
    normalized = _normalize_state(state)
    normalized["updated_at"] = _utc_now()
    path = state_file_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(normalized, indent=2, sort_keys=True), encoding="utf-8")


def resolve_optimizer_mode(requested: Optional[str]) -> str:
    mode = str(requested or "auto").strip().lower()
    env_mode = str(os.environ.get("OSMOSIS_OPTIMIZER", "")).strip().lower()
    if env_mode in OPTIMIZER_MODES:
        mode = env_mode
    if mode not in OPTIMIZER_MODES:
        raise ValueError(
            f"Unsupported optimizer mode '{mode}'. Expected one of: {sorted(OPTIMIZER_MODES)}"
        )
    return mode


def resolve_subsystem_mode(
    requested: Optional[str],
    *,
    subsystem_env_var: Optional[str] = None,
) -> str:
    global_mode = resolve_optimizer_mode("auto")
    if global_mode == "off":
        return "off"

    mode = str(requested or "auto").strip().lower()
    if subsystem_env_var:
        env_mode = str(os.environ.get(subsystem_env_var, "")).strip().lower()
        if env_mode in OPTIMIZER_MODES:
            mode = env_mode
    if mode not in OPTIMIZER_MODES:
        raise ValueError(
            f"Unsupported mode '{mode}'. Expected one of: {sorted(OPTIMIZER_MODES)}"
        )
    return mode


def record_subsystem_selection(
    *,
    subsystem: str,
    profile_id: str,
    pack_version: str,
    selector_version: str,
    selection_latency_ms: float,
    quality_ok: bool,
    fallback_reason: str = "",
    quality_score: Optional[float] = None,
) -> Dict[str, Any]:
    name = str(subsystem or "").strip().lower() or "unknown"
    pid = str(profile_id or "").strip() or "unknown"
    state = load_state()
    subsystems = state.setdefault("subsystems", {})
    entry = subsystems.setdefault(name, _subsystem_default())

    entry["runs"] = int(entry.get("runs", 0) or 0) + 1
    entry["active_profile_id"] = pid
    entry["pack_version"] = str(pack_version or "")
    entry["selector_version"] = str(selector_version or "")
    entry["last_selection_latency_ms"] = round(float(selection_latency_ms or 0.0), 4)

    candidates = entry.setdefault("candidates", {})
    candidates[pid] = int(candidates.get(pid, 0) or 0) + 1

    if quality_score is not None:
        scores = entry.setdefault("scores", {})
        prev = float(scores.get(pid, 0.0) or 0.0)
        scores[pid] = round(0.8 * prev + 0.2 * float(quality_score), 6)

    if not quality_ok:
        entry["canary_status"] = "rollback"
        entry["rollback_reason"] = str(fallback_reason or "quality_regression")
        entry["canary_passes"] = 0
        state["rollback_marker"] = {
            "subsystem": name,
            "profile_id": pid,
            "reason": entry["rollback_reason"],
            "at": _utc_now(),
        }
    else:
        current = str(entry.get("canary_status", "shadow"))
        promoted = str(entry.get("promoted_profile_id", ""))
        if promoted == pid:
            entry["canary_status"] = "promoted"
            entry["rollback_reason"] = ""
            entry["canary_passes"] = int(entry.get("canary_passes", 0) or 0) + 1
        else:
            passes = int(entry.get("canary_passes", 0) or 0) + 1
            entry["canary_passes"] = passes
            if current in {"shadow", "rollback"}:
                entry["canary_status"] = "canary"
            elif current == "canary" and passes >= 2:
                entry["canary_status"] = "promoted"
                entry["promoted_profile_id"] = pid
                entry["rollback_reason"] = ""

    save_state(state)
    return get_subsystem_metadata(name)


def get_subsystem_metadata(subsystem: str) -> Dict[str, Any]:
    name = str(subsystem or "").strip().lower() or "unknown"
    state = load_state()
    entry = dict(state.get("subsystems", {}).get(name, _subsystem_default()))
    return {
        "canary_status": str(entry.get("canary_status", "shadow")),
        "promoted_profile_id": str(entry.get("promoted_profile_id", "")),
        "active_profile_id": str(entry.get("active_profile_id", "")),
        "rollback_reason": str(entry.get("rollback_reason", "")),
        "pack_version": str(entry.get("pack_version", "")),
        "selector_version": str(entry.get("selector_version", "")),
        "selection_latency_ms": round(float(entry.get("last_selection_latency_ms", 0.0) or 0.0), 4),
        "runs": int(entry.get("runs", 0) or 0),
        "shadow_runs": int(entry.get("shadow_runs", 0) or 0),
        "canary_attempts": int(entry.get("canary_attempts", 0) or 0),
        "canary_passes": int(entry.get("canary_passes", 0) or 0),
    }


def get_optimizer_status() -> Dict[str, Any]:
    state = load_state()
    subsystems_raw = state.get("subsystems", {})
    status_subsystems: Dict[str, Any] = {}
    if isinstance(subsystems_raw, dict):
        for name, _ in sorted(subsystems_raw.items(), key=lambda item: str(item[0])):
            status_subsystems[str(name)] = get_subsystem_metadata(str(name))
    return {
        "schema_version": STATE_SCHEMA_VERSION,
        "optimizer_mode": resolve_optimizer_mode("auto"),
        "state_path": str(state_file_path()),
        "updated_at": str(state.get("updated_at", "")),
        "rollback_marker": dict(state.get("rollback_marker", {}) or {}),
        "subsystems": status_subsystems,
    }


def reset_optimizer_state() -> Dict[str, Any]:
    home = optimizer_home()
    if home.exists():
        shutil.rmtree(home)
    return get_optimizer_status()
