"""Deterministic compression-profile selector backed by local optimizer state."""

from __future__ import annotations

import hashlib
import json
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np

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


COMPRESSION_SELECTOR_VERSION = "compression-selector-v1"
COMPRESSION_ENGINE_MODES = {"auto", "off"}
COMPRESSION_BUDGET_MS = 5000
DEFAULT_COMPRESSION_PACK_VERSION = "builtin-compression-pack-v1"


@dataclass(frozen=True)
class CompressionSelection:
    """Selected compression profile metadata."""

    compression_engine_mode: str
    compression_profile: str
    profile_id: str
    pack_version: str
    selector_version: str
    zstd_levels: Tuple[int, ...]
    selection_latency_ms: float
    fallback_reason: str
    canary_status: str
    rollback_reason: str

    def to_report_dict(self) -> Dict[str, Any]:
        return {
            "optimizer_mode": self.compression_engine_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,
            "zstd_levels": list(self.zstd_levels),
        }


def resolve_compression_engine_mode(requested: Optional[str]) -> str:
    return resolve_subsystem_mode(
        requested,
        subsystem_env_var="OSMOSIS_COMPRESSION_ENGINE",
    )


def _canonical_json_bytes(data: Dict[str, Any]) -> bytes:
    return json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8")


def _sha256_json(data: Dict[str, Any]) -> str:
    return hashlib.sha256(_canonical_json_bytes(data)).hexdigest()


def _builtin_pack() -> Dict[str, Any]:
    profiles = [
        {
            "profile_id": "lossless_fast_v1",
            "compression_profile": "lossless-max",
            "zstd_levels": [3, 6, 9],
        },
        {
            "profile_id": "lossless_balanced_v1",
            "compression_profile": "lossless-max",
            "zstd_levels": [3, 9, 19],
        },
        {
            "profile_id": "lossless_maxratio_v1",
            "compression_profile": "lossless-max",
            "zstd_levels": [9, 19, 22],
        },
    ]
    payload = {
        "pack_version": DEFAULT_COMPRESSION_PACK_VERSION,
        "selector_version": COMPRESSION_SELECTOR_VERSION,
        "profiles": profiles,
    }
    payload["checksum_sha256"] = _sha256_json(payload)
    return payload


def _pack_file_path() -> Optional[Path]:
    pack_env = str(os.environ.get("OSMOSIS_COMPRESSION_PACK", "")).strip()
    if pack_env:
        return Path(pack_env)
    candidate = optimizer_home() / "packs" / "compression_pack.json"
    return candidate if candidate.exists() else None


def _load_pack() -> Tuple[Dict[str, Any], str]:
    path = _pack_file_path()
    if path is None:
        return _builtin_pack(), "builtin_pack"
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return _builtin_pack(), "invalid_pack_json_fallback_builtin"
    if not isinstance(data, dict):
        return _builtin_pack(), "invalid_pack_schema_fallback_builtin"

    payload = {
        "pack_version": data.get("pack_version", ""),
        "selector_version": data.get("selector_version", ""),
        "profiles": data.get("profiles", []),
    }
    expected = str(data.get("checksum_sha256", "")).strip().lower()
    computed = hashlib.sha256(_canonical_json_bytes(payload)).hexdigest()
    if expected and expected != computed:
        return _builtin_pack(), "pack_checksum_mismatch_fallback_builtin"
    if not expected:
        data["checksum_sha256"] = computed

    trust_block = data.get("trust")
    if isinstance(trust_block, dict):
        try:
            from model_osmosis.format.trust import verify_trust_block_with_backend

            verify = verify_trust_block_with_backend(
                trust_block=trust_block,
                payload_bytes=_canonical_json_bytes(payload),
                backend_name=str(trust_block.get("backend") or ""),
            )
            status = str(verify.get("trust_status", "unsigned"))
            if status not in {"trusted", "signed_untrusted"}:
                return _builtin_pack(), "pack_signature_invalid_fallback_builtin"
        except Exception:
            return _builtin_pack(), "pack_signature_invalid_fallback_builtin"

    return data, "external_pack"


def _profiles_from_pack(pack: Dict[str, Any]) -> List[Dict[str, Any]]:
    profiles_raw = pack.get("profiles", [])
    if not isinstance(profiles_raw, list):
        return []
    profiles: List[Dict[str, Any]] = []
    for item in profiles_raw:
        if not isinstance(item, dict):
            continue
        profile_id = str(item.get("profile_id", "")).strip()
        comp_profile = str(item.get("compression_profile", "")).strip()
        levels = item.get("zstd_levels", [])
        if not profile_id or comp_profile != "lossless-max" or not isinstance(levels, list):
            continue
        try:
            level_tuple = tuple(int(level) for level in levels)
        except Exception:
            continue
        if not level_tuple:
            continue
        profiles.append(
            {
                "profile_id": profile_id,
                "compression_profile": comp_profile,
                "zstd_levels": level_tuple,
            }
        )
    return sorted(profiles, key=lambda item: item["profile_id"])


def _estimate_raw_size_mb(osmo_model: Any) -> float:
    raw_bytes = 0
    params = getattr(osmo_model, "parameters", {}) or {}
    for arr in params.values():
        raw_bytes += int(np.asarray(arr).nbytes)
    knowledge = getattr(osmo_model, "knowledge", {}) or {}
    for arr in knowledge.values():
        raw_bytes += int(np.asarray(arr).nbytes)
    return float(raw_bytes) / (1024.0 * 1024.0)


def _score_profile(profile: Dict[str, Any], *, raw_size_mb: float) -> float:
    levels = tuple(profile.get("zstd_levels", (3, 9, 19)))
    avg_level = float(sum(levels) / max(len(levels), 1))
    if raw_size_mb < 32.0:
        target = 6.0
    elif raw_size_mb < 256.0:
        target = 10.0
    else:
        target = 16.0
    return -abs(avg_level - target)


def _default_selection(
    *,
    mode: str,
    compression_profile: str,
    fallback_reason: str,
    profile_id: str,
    pack_version: str = DEFAULT_COMPRESSION_PACK_VERSION,
) -> CompressionSelection:
    metadata = get_subsystem_metadata("compression")
    levels = (3, 9, 19) if compression_profile == "lossless-max" else ()
    return CompressionSelection(
        compression_engine_mode=mode,
        compression_profile=compression_profile,
        profile_id=profile_id,
        pack_version=pack_version,
        selector_version=COMPRESSION_SELECTOR_VERSION,
        zstd_levels=levels,
        selection_latency_ms=0.0,
        fallback_reason=fallback_reason,
        canary_status=str(metadata.get("canary_status", "shadow")),
        rollback_reason=str(metadata.get("rollback_reason", "")),
    )


def select_compression_selection(
    *,
    osmo_model: Any,
    compression_profile: str,
    compression_engine_mode: str = "auto",
    budget_ms: int = COMPRESSION_BUDGET_MS,
    update_state: bool = True,
) -> CompressionSelection:
    """Select compression tuning profile deterministically within lossless constraints."""
    mode = resolve_compression_engine_mode(compression_engine_mode)
    normalized_profile = str(compression_profile or "legacy").strip().lower()
    if normalized_profile != "lossless-max":
        return _default_selection(
            mode=mode,
            compression_profile=normalized_profile,
            fallback_reason="non_lossless_profile_passthrough",
            profile_id="legacy_passthrough",
            pack_version=DEFAULT_COMPRESSION_PACK_VERSION,
        )
    if mode == "off":
        return _default_selection(
            mode=mode,
            compression_profile=normalized_profile,
            fallback_reason="compression_engine_off",
            profile_id="lossless_balanced_v1",
            pack_version=DEFAULT_COMPRESSION_PACK_VERSION,
        )

    started = time.perf_counter()
    pack, pack_reason = _load_pack()
    profiles = _profiles_from_pack(pack)
    if not profiles:
        return _default_selection(
            mode=mode,
            compression_profile=normalized_profile,
            fallback_reason="no_valid_profiles",
            profile_id="lossless_balanced_v1",
            pack_version=str(pack.get("pack_version") or DEFAULT_COMPRESSION_PACK_VERSION),
        )

    raw_size_mb = _estimate_raw_size_mb(osmo_model)
    scored = sorted(
        (
            (
                _score_profile(profile, raw_size_mb=raw_size_mb),
                str(profile["profile_id"]),
                profile,
            )
            for profile in profiles
        ),
        key=lambda item: (-float(item[0]), item[1]),
    )
    selected = scored[0][2]
    elapsed_ms = (time.perf_counter() - started) * 1000.0
    if elapsed_ms > float(budget_ms):
        return _default_selection(
            mode=mode,
            compression_profile=normalized_profile,
            fallback_reason="selection_budget_exceeded",
            profile_id="lossless_balanced_v1",
            pack_version=str(pack.get("pack_version") or DEFAULT_COMPRESSION_PACK_VERSION),
        )

    if update_state:
        metadata = record_subsystem_selection(
            subsystem="compression",
            profile_id=str(selected["profile_id"]),
            pack_version=str(pack.get("pack_version") or DEFAULT_COMPRESSION_PACK_VERSION),
            selector_version=str(pack.get("selector_version") or COMPRESSION_SELECTOR_VERSION),
            selection_latency_ms=elapsed_ms,
            quality_ok=True,  # lossless profile tuning preserves model fidelity by design
            fallback_reason=pack_reason,
            quality_score=float(scored[0][0]),
        )
    else:
        metadata = get_subsystem_metadata("compression")

    return CompressionSelection(
        compression_engine_mode=mode,
        compression_profile=normalized_profile,
        profile_id=str(selected["profile_id"]),
        pack_version=str(pack.get("pack_version") or DEFAULT_COMPRESSION_PACK_VERSION),
        selector_version=str(pack.get("selector_version") or COMPRESSION_SELECTOR_VERSION),
        zstd_levels=tuple(int(level) for level in selected.get("zstd_levels", (3, 9, 19))),
        selection_latency_ms=elapsed_ms,
        fallback_reason=pack_reason,
        canary_status=str(metadata.get("canary_status", "shadow")),
        rollback_reason=str(metadata.get("rollback_reason", "")),
    )
