"""Deterministic weight-only merge backend across converted Osmo models."""

from __future__ import annotations

import copy
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np

from model_osmosis.format.converters import UniversalConverter
from model_osmosis.format.osmo_format import OsmoModel
from model_osmosis.merging.smart_merge_engine import (
    collect_tensor_groups_from_osmo_parameters,
    default_selection,
    select_recipe_for_tensor_groups,
)
from model_osmosis.merging.strategies import get_strategy


@dataclass
class WeightMergeResult:
    """Result of a weight-only merge."""

    merged_osmo: OsmoModel
    report: Dict[str, Any]


def _iso_utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _normalize_weights(weights: Optional[Sequence[float]], model_count: int) -> Optional[List[float]]:
    if weights is None:
        return None
    values = [float(w) for w in weights]
    if len(values) != model_count:
        raise ValueError(
            f"Weight-merge expected {model_count} weights, got {len(values)}."
        )
    total = float(sum(values))
    if total <= 0:
        raise ValueError("Weight-merge weights must sum to a positive value.")
    return [v / total for v in values]


def _to_osmo_models(
    models: Sequence[Any],
    *,
    source_refs: Sequence[str],
    allow_lossy_cast: bool = False,
) -> List[OsmoModel]:
    converter = UniversalConverter()
    converted: List[OsmoModel] = []
    for idx, model in enumerate(models):
        ref = source_refs[idx] if idx < len(source_refs) else f"model_{idx + 1}"
        if isinstance(model, OsmoModel):
            osmo = model
        else:
            osmo = converter.to_osmo(
                model,
                name=ref,
                allow_lossy_cast=allow_lossy_cast,
            )
        converted.append(osmo)
    return converted


def _resolve_strategy(strategy: str, *, seed: int) -> Tuple[str, Any]:
    requested = str(strategy or "weighted_average").strip().lower()
    used = "weighted_average" if requested == "auto" else requested
    kwargs: Dict[str, Any] = {}
    if used == "dare":
        kwargs["seed"] = int(seed)
    return used, get_strategy(used, **kwargs)


def _analyze_compatibility(
    osmos: Sequence[OsmoModel],
) -> Tuple[List[str], Dict[str, int]]:
    union_names = sorted({name for osmo in osmos for name in osmo.parameters.keys()})
    reasons = {
        "missing": 0,
        "shape_mismatch": 0,
        "dtype_mismatch": 0,
    }
    compatible: List[str] = []

    for name in union_names:
        tensors: List[np.ndarray] = []
        missing = False
        for osmo in osmos:
            if name not in osmo.parameters:
                missing = True
                break
            tensors.append(osmo.parameters[name])
        if missing:
            reasons["missing"] += 1
            continue

        first_shape = tuple(tensors[0].shape)
        if any(tuple(t.shape) != first_shape for t in tensors[1:]):
            reasons["shape_mismatch"] += 1
            continue

        first_dtype = str(tensors[0].dtype)
        if any(str(t.dtype) != first_dtype for t in tensors[1:]):
            reasons["dtype_mismatch"] += 1
            continue

        compatible.append(name)

    return compatible, reasons


def _coverage_by_bytes(
    osmos: Sequence[OsmoModel],
    compatible_names: Sequence[str],
    source_refs: Sequence[str],
) -> Tuple[Dict[str, float], float]:
    coverage: Dict[str, float] = {}
    for idx, osmo in enumerate(osmos):
        label = source_refs[idx] if idx < len(source_refs) else (osmo.manifest.name or f"model_{idx + 1}")
        total = float(sum(arr.nbytes for arr in osmo.parameters.values()))
        merged = float(sum(osmo.parameters[name].nbytes for name in compatible_names if name in osmo.parameters))
        ratio = 0.0 if total <= 0 else merged / total
        coverage[str(label)] = float(ratio)
    min_ratio = min(coverage.values()) if coverage else 0.0
    return coverage, float(min_ratio)


def _build_merged_template(
    osmos: Sequence[OsmoModel],
    merged_params: Dict[str, np.ndarray],
    *,
    strategy_used: str,
    source_refs: Sequence[str],
) -> OsmoModel:
    template = copy.deepcopy(osmos[0])
    template.parameters = {name: merged_params[name] for name in sorted(merged_params.keys())}
    template.architecture.layers = [
        layer for layer in template.architecture.layers if layer.name in template.parameters
    ]
    template.architecture.num_layers = len(template.architecture.layers)

    source_list = [str(r) for r in source_refs]
    source_text = ", ".join(source_list[:3])
    if len(source_list) > 3:
        source_text += ", ..."

    base_name = template.manifest.name or "model"
    template.manifest.name = f"{base_name}-weight-merged"
    template.manifest.description = (
        f"Weight-merged ({strategy_used}) from {source_text or 'multiple sources'}."
    )
    template.manifest.parameter_count = int(sum(arr.size for arr in template.parameters.values()))
    template.manifest.created_at = _iso_utc_now()
    tags = list(template.manifest.tags or [])
    tags.extend(["weight-merge", strategy_used])
    template.manifest.tags = sorted(set(str(t) for t in tags))
    template.fingerprint = template._compute_fingerprint()
    return template


def merge_weights_to_osmo(
    models: Sequence[Any],
    *,
    source_refs: Sequence[str],
    strategy: str = "weighted_average",
    weights: Optional[Sequence[float]] = None,
    coverage_floor: float = 0.70,
    seed: int = 42,
    allow_lossy_cast: bool = False,
    strategy_engine_mode: str = "auto",
) -> WeightMergeResult:
    """Merge only exact-compatible tensors across models and return an OsmoModel."""
    if len(models) < 2:
        raise ValueError("Weight-merge requires at least two models.")
    floor = float(coverage_floor)
    if floor < 0.0 or floor > 1.0:
        raise ValueError("coverage_floor must be within [0.0, 1.0].")

    osmos = _to_osmo_models(
        models,
        source_refs=source_refs,
        allow_lossy_cast=allow_lossy_cast,
    )
    normalized_weights = _normalize_weights(weights, len(osmos))
    compatible_names, reason_counts = _analyze_compatibility(osmos)
    selection = default_selection(
        strategy_name="weighted_average",
        strategy_engine_mode="off",
        fallback_reason="explicit_strategy",
    )
    strategy_requested = str(strategy)
    strategy_requested_norm = strategy_requested.strip().lower()
    if strategy_requested_norm == "auto":
        tensor_groups = collect_tensor_groups_from_osmo_parameters(
            osmos,
            compatible_names,
            seed=seed,
        )
        selection = select_recipe_for_tensor_groups(
            tensor_groups=tensor_groups,
            model_count=len(osmos),
            seed=seed,
            strategy_engine_mode=strategy_engine_mode,
            user_weights_provided=normalized_weights is not None,
            context_label="weight-merge",
            source_refs=list(source_refs),
        )
        strategy_used = selection.strategy_name
        strategy_kwargs = dict(selection.strategy_kwargs)
        if strategy_used == "dare" and "seed" not in strategy_kwargs:
            strategy_kwargs["seed"] = int(seed)
        strategy_impl = get_strategy(strategy_used, **strategy_kwargs)
    else:
        strategy_used, strategy_impl = _resolve_strategy(strategy, seed=seed)
        selection = default_selection(
            strategy_name=strategy_used,
            strategy_engine_mode="off",
            recipe_id=f"explicit::{strategy_used}",
            fallback_reason="explicit_strategy",
        )

    merged_params: Dict[str, np.ndarray] = {}
    for name in compatible_names:
        tensors = [osmo.parameters[name] for osmo in osmos]
        merged_params[name] = strategy_impl.merge_tensors(
            tensors,
            weights=normalized_weights,
        )

    coverage_map, min_coverage = _coverage_by_bytes(osmos, compatible_names, source_refs)
    compatible_count = len(compatible_names)
    skipped_count = int(sum(reason_counts.values()))
    merged_bytes = int(sum(arr.nbytes for arr in merged_params.values()))

    report = {
        "mode": "weight-merge",
        "strategy_requested": strategy_requested,
        "strategy_used": strategy_used,
        "strategy_kwargs": dict(sorted((selection.strategy_kwargs or {}).items(), key=lambda item: item[0])),
        "model_count": len(osmos),
        "compatible_tensor_count": compatible_count,
        "skipped_tensor_count": skipped_count,
        "merged_bytes": merged_bytes,
        "coverage_ratio_by_bytes": {
            str(k): float(v) for k, v in sorted(coverage_map.items(), key=lambda item: item[0])
        },
        "min_coverage_ratio": float(min_coverage),
        "coverage_floor": floor,
        "skip_reason_counts": {
            "dtype_mismatch": int(reason_counts["dtype_mismatch"]),
            "missing": int(reason_counts["missing"]),
            "shape_mismatch": int(reason_counts["shape_mismatch"]),
        },
        "strategy_engine_mode": selection.strategy_engine_mode,
        "recipe_id": selection.recipe_id,
        "recipe_pack_version": selection.recipe_pack_version,
        "selector_version": selection.selector_version,
        "selection_latency_ms": selection.selection_latency_ms,
        "fallback_reason": selection.fallback_reason,
        "optimizer_mode": selection.strategy_engine_mode,
        "profile_id": selection.recipe_id,
        "pack_version": selection.recipe_pack_version,
        "canary_status": selection.canary_status,
        "rollback_reason": selection.rollback_reason,
    }

    if min_coverage < floor:
        raise RuntimeError(
            "Weight-merge coverage floor not met: "
            f"min_coverage_ratio={min_coverage:.4f} < coverage_floor={floor:.4f}. "
            f"compatible_tensors={compatible_count}, skipped_tensors={skipped_count}."
        )

    merged_osmo = _build_merged_template(
        osmos,
        merged_params,
        strategy_used=strategy_used,
        source_refs=source_refs,
    )
    return WeightMergeResult(merged_osmo=merged_osmo, report=report)
