"""Data Mixing Optimizer (v0.48.0 Part B — BETA).

Run short proxy-training runs with different per-dataset mixture weights and
fit a Gaussian Process surrogate to recommend the optimal mixture.

This module ships the schema + budget accountant + recipe writer. The live
Bayesian optimisation loop (``scikit-optimize``) is wired through a
runtime-injected ``OptimizerProtocol`` so unit tests can drive deterministic
mock optimisers and the library import is lazy (matches the project's
``[optional-extras]`` policy — heavy deps never crash ``soup data --help``).

CLI surface:
    soup data mix --optimize --budget 1h --datasets a.jsonl,b.jsonl,c.jsonl
    soup data mix --apply <recipe.yaml>

Security:
- All input/output paths are containment-checked via ``utils.paths.is_under_cwd``.
- ``--budget`` is wall-clock capped; partial results returned on early exit.
- Dataset paths reject null bytes / oversize / non-string.
- ``scikit-optimize`` is lazy-imported; missing dep surfaces a friendly advisory.
"""

from __future__ import annotations

import json
import math
import os
import stat
import tempfile
import time
from dataclasses import dataclass
from typing import Callable, List, Mapping, Optional, Protocol, Sequence, Tuple

# --- Limits / constants ---------------------------------------------------

_MAX_DATASETS = 32  # mirrors v0.42.0 interleave cap
_MAX_PROBES = 256  # hard ceiling on N short proxy runs
_DEFAULT_PROBES = 8
_MIN_BUDGET_SECONDS = 60  # 1 minute
_MAX_BUDGET_SECONDS = 24 * 60 * 60  # 24 hours
_MAX_PATH_LEN = 4096
_MAX_RECIPE_BYTES = 256 * 1024  # mirror v0.39.0 Part E
_MAX_LOSS = 1e6
_FLOAT_TOL = 1e-6

__all__ = [
    "MixCandidate",
    "MixOptimizationReport",
    "MixOptimizationPlan",
    "BudgetTracker",
    "OptimizerProtocol",
    "validate_datasets",
    "parse_budget",
    "build_optimization_plan",
    "render_mix_recipe_yaml",
    "write_mix_recipe",
    "run_mix_optimizer",
]


# --- Dataclasses ----------------------------------------------------------


@dataclass(frozen=True)
class MixCandidate:
    """One proxy-run candidate: mixture weights + observed eval loss.

    Attributes:
        weights: Per-dataset weights summing to 1.0 ± 1e-6.
        eval_loss: Observed eval loss after the short proxy run.
        wall_clock_seconds: Time spent on this candidate.
    """

    weights: Tuple[float, ...]
    eval_loss: float
    wall_clock_seconds: float

    def __post_init__(self) -> None:
        if isinstance(self.weights, bool) or not isinstance(
            self.weights, tuple
        ):
            raise TypeError(
                f"weights must be tuple, got {type(self.weights).__name__}"
            )
        if not self.weights:
            raise ValueError("weights must be non-empty")
        for w in self.weights:
            if isinstance(w, bool):
                raise ValueError("weight must be float, not bool")
            if not isinstance(w, (int, float)):
                raise TypeError(f"weight must be float, got {type(w).__name__}")
            fw = float(w)
            if not math.isfinite(fw):
                raise ValueError(f"weight must be finite (got {w!r})")
            if fw < 0.0 or fw > 1.0:
                raise ValueError(f"weight must be in [0, 1], got {fw}")
        total = sum(float(w) for w in self.weights)
        if abs(total - 1.0) > _FLOAT_TOL:
            raise ValueError(
                f"weights must sum to 1.0 ± {_FLOAT_TOL} (got {total})"
            )
        for name, value in (
            ("eval_loss", self.eval_loss),
            ("wall_clock_seconds", self.wall_clock_seconds),
        ):
            if isinstance(value, bool):
                raise ValueError(f"{name} must be float, not bool")
            if not isinstance(value, (int, float)):
                raise TypeError(
                    f"{name} must be float, got {type(value).__name__}"
                )
            fv = float(value)
            if not math.isfinite(fv):
                raise ValueError(f"{name} must be finite (got {value!r})")
            if fv < 0.0:
                raise ValueError(f"{name} must be >= 0 (got {fv})")
        if self.eval_loss > _MAX_LOSS:
            raise ValueError(
                f"eval_loss exceeds sanity cap {_MAX_LOSS} (got {self.eval_loss})"
            )


@dataclass(frozen=True)
class MixOptimizationReport:
    """Result of a mixing-optimizer run.

    Attributes:
        datasets: The dataset paths in canonical order.
        candidates: Tuple of evaluated candidates (chronological).
        best_weights: Mixture with the lowest eval loss observed.
        best_eval_loss: The corresponding loss.
        partial: True when the budget tripped before all candidates ran.
        elapsed_seconds: Total wall-clock spent (sum of per-candidate time).
    """

    datasets: Tuple[str, ...]
    candidates: Tuple[MixCandidate, ...]
    best_weights: Tuple[float, ...]
    best_eval_loss: float
    partial: bool
    elapsed_seconds: float


@dataclass(frozen=True)
class MixOptimizationPlan:
    """Validated plan for a mixing-optimization invocation.

    Attributes:
        datasets: Canonical dataset paths (real-paths within cwd).
        num_probes: How many proxy runs to attempt.
        budget_seconds: Hard wall-clock cap.
        seed: RNG seed for the optimizer.
    """

    datasets: Tuple[str, ...]
    num_probes: int
    budget_seconds: int
    seed: int


# --- Validation helpers ---------------------------------------------------


def _reject_bool_int(name: str, value) -> int:
    if isinstance(value, bool):
        raise ValueError(f"{name} must be int, not bool")
    if not isinstance(value, int):
        raise TypeError(f"{name} must be int, got {type(value).__name__}")
    return value


def _check_str_path(name: str, value) -> str:
    if not isinstance(value, str):
        raise TypeError(f"{name} must be str, got {type(value).__name__}")
    if not value:
        raise ValueError(f"{name} must be non-empty")
    if "\x00" in value:
        raise ValueError(f"{name} must not contain null bytes")
    if len(value) > _MAX_PATH_LEN:
        raise ValueError(
            f"{name} length {len(value)} exceeds cap {_MAX_PATH_LEN}"
        )
    return value


def validate_datasets(raw: Sequence[str]) -> Tuple[str, ...]:
    """Validate dataset paths: containment + dedup + bounds.

    Args:
        raw: Sequence of dataset paths (relative or absolute).

    Returns:
        Tuple of real-path strings, all confined to cwd.

    Raises:
        TypeError / ValueError on bad input.
    """
    from soup_cli.utils.paths import is_under_cwd

    if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)):
        raise TypeError("datasets must be a non-string Sequence")
    if len(raw) < 2:
        raise ValueError(
            f"datasets must contain at least 2 entries (got {len(raw)})"
        )
    if len(raw) > _MAX_DATASETS:
        raise ValueError(
            f"datasets has {len(raw)} entries; cap is {_MAX_DATASETS}"
        )
    seen: List[str] = []
    for item in raw:
        path = _check_str_path("dataset", item)
        # Symlink check on the ORIGINAL path BEFORE realpath (which would
        # follow the symlink, defeating the check). Matches v0.46.0 Part A
        # `_reject_symlink_target` pattern.
        try:
            st = os.lstat(path)
        except FileNotFoundError:
            st = None
        except OSError as exc:
            raise ValueError(
                f"dataset path is not stat-able: {os.path.basename(path)!r}"
            ) from exc
        if st is not None and stat.S_ISLNK(st.st_mode):
            raise ValueError(
                f"dataset path is a symlink (rejected for safety): "
                f"{os.path.basename(path)!r}"
            )
        real = os.path.realpath(path)
        if not is_under_cwd(real):
            raise ValueError(
                f"dataset path is outside cwd: {os.path.basename(real)!r}"
            )
        if real in seen:
            raise ValueError(
                f"duplicate dataset path: {os.path.basename(real)!r}"
            )
        seen.append(real)
    if len(seen) < 2:
        raise ValueError(
            "data mix requires at least 2 distinct datasets"
        )
    return tuple(seen)


def parse_budget(raw: str) -> int:
    """Parse a wall-clock budget string into seconds.

    Accepts:
        ``30s`` / ``5m`` / ``1h`` / ``600`` (bare seconds).

    Raises:
        ValueError on invalid format / out-of-bounds.
    """
    if not isinstance(raw, str):
        raise TypeError(f"budget must be str, got {type(raw).__name__}")
    s = raw.strip().lower()
    if not s:
        raise ValueError("budget must be non-empty")
    if "\x00" in s:
        raise ValueError("budget must not contain null bytes")

    multiplier = 1
    body = s
    if s.endswith("s"):
        body = s[:-1]
    elif s.endswith("m"):
        body = s[:-1]
        multiplier = 60
    elif s.endswith("h"):
        body = s[:-1]
        multiplier = 3600
    if not body or not body.isdigit():
        raise ValueError(
            f"budget must be digits + optional suffix (s/m/h), got {raw!r}"
        )
    seconds = int(body) * multiplier
    if seconds < _MIN_BUDGET_SECONDS or seconds > _MAX_BUDGET_SECONDS:
        raise ValueError(
            f"budget must resolve to [{_MIN_BUDGET_SECONDS}, "
            f"{_MAX_BUDGET_SECONDS}] seconds (got {seconds})"
        )
    return seconds


def build_optimization_plan(
    datasets: Sequence[str],
    *,
    budget: str = "1h",
    num_probes: int = _DEFAULT_PROBES,
    seed: int = 42,
) -> MixOptimizationPlan:
    """Validate args and produce a frozen plan."""
    ds = validate_datasets(datasets)
    budget_seconds = parse_budget(budget)
    nb = _reject_bool_int("num_probes", num_probes)
    if nb < 1 or nb > _MAX_PROBES:
        raise ValueError(
            f"num_probes must be in [1, {_MAX_PROBES}], got {nb}"
        )
    sd = _reject_bool_int("seed", seed)
    if sd < 0 or sd > 2**31 - 1:
        raise ValueError(f"seed must be in [0, 2**31-1], got {sd}")
    return MixOptimizationPlan(
        datasets=ds,
        num_probes=nb,
        budget_seconds=budget_seconds,
        seed=sd,
    )


# --- Budget tracker -------------------------------------------------------


class BudgetTracker:
    """Wall-clock budget accountant.

    Used by :func:`run_mix_optimizer` to terminate the BO loop when the
    cumulative time exceeds the configured budget. Partial results are
    surfaced via :class:`MixOptimizationReport` with ``partial=True``.
    """

    def __init__(
        self,
        budget_seconds: int,
        *,
        clock: Optional[Callable[[], float]] = None,
    ) -> None:
        bs = _reject_bool_int("budget_seconds", budget_seconds)
        if bs < _MIN_BUDGET_SECONDS or bs > _MAX_BUDGET_SECONDS:
            raise ValueError(
                f"budget_seconds must be in "
                f"[{_MIN_BUDGET_SECONDS}, {_MAX_BUDGET_SECONDS}], got {bs}"
            )
        self._budget = bs
        self._clock = clock or time.monotonic
        self._started: Optional[float] = None

    def start(self) -> None:
        if self._started is not None:
            raise RuntimeError("BudgetTracker.start called twice")
        self._started = self._clock()

    @property
    def elapsed(self) -> float:
        if self._started is None:
            return 0.0
        return max(0.0, self._clock() - self._started)

    @property
    def remaining(self) -> float:
        return max(0.0, self._budget - self.elapsed)

    def exceeded(self) -> bool:
        return self.elapsed >= self._budget


# --- Optimizer protocol ---------------------------------------------------


class OptimizerProtocol(Protocol):
    """Duck-typed interface for the BO backend.

    Implementations must produce non-negative weights that sum to 1.0; the
    runner re-normalises to defend against floating-point drift.
    """

    def ask(self) -> Tuple[float, ...]:
        """Return the next candidate weights."""

    def tell(self, weights: Tuple[float, ...], loss: float) -> None:
        """Record an observation."""


def _build_default_optimizer(
    num_datasets: int, seed: int
) -> OptimizerProtocol:
    """Return a deterministic Dirichlet-like sampler when scikit-optimize is
    not installed. Otherwise wrap ``skopt.Optimizer`` (lazy import).
    """
    import random

    rng = random.Random(seed)

    class _Dirichlet:
        def ask(self) -> Tuple[float, ...]:
            # Symmetric Dirichlet(α=1) via independent exponentials.
            raw = [rng.expovariate(1.0) for _ in range(num_datasets)]
            total = sum(raw) or 1.0
            return tuple(r / total for r in raw)

        def tell(self, weights: Tuple[float, ...], loss: float) -> None:
            return

    return _Dirichlet()


def _renormalize(weights: Sequence[float]) -> Tuple[float, ...]:
    """Clip + renormalise to a valid simplex point."""
    clipped = [max(0.0, float(w)) for w in weights]
    total = sum(clipped)
    if total <= 0.0:
        n = len(clipped)
        return tuple([1.0 / n] * n) if n else ()
    return tuple(c / total for c in clipped)


# --- Optimizer runner -----------------------------------------------------


def run_mix_optimizer(
    plan: MixOptimizationPlan,
    proxy_run: Callable[[Tuple[float, ...]], float],
    *,
    optimizer: Optional[OptimizerProtocol] = None,
    clock: Optional[Callable[[], float]] = None,
) -> MixOptimizationReport:
    """Run the BO loop over the validated plan.

    Args:
        plan: A :class:`MixOptimizationPlan` from
            :func:`build_optimization_plan`.
        proxy_run: Callable that takes weights and returns observed eval loss.
            In the live wiring this calls a short ``soup train`` invocation;
            in tests it is mocked.
        optimizer: Optional injected :class:`OptimizerProtocol`. Defaults to
            the Dirichlet sampler when absent.
        clock: Optional monotonic clock callable (testability).

    Returns:
        A :class:`MixOptimizationReport`. ``partial=True`` when budget tripped.
    """
    if not isinstance(plan, MixOptimizationPlan):
        raise TypeError(
            f"plan must be MixOptimizationPlan, got {type(plan).__name__}"
        )
    if not callable(proxy_run):
        raise TypeError("proxy_run must be callable")
    if optimizer is not None and (
        not hasattr(optimizer, "ask") or not hasattr(optimizer, "tell")
    ):
        raise TypeError(
            "optimizer must implement OptimizerProtocol (ask + tell)"
        )

    opt = optimizer or _build_default_optimizer(
        len(plan.datasets), plan.seed
    )
    tracker = BudgetTracker(plan.budget_seconds, clock=clock)
    tracker.start()

    candidates: List[MixCandidate] = []
    best_weights: Optional[Tuple[float, ...]] = None
    best_loss: float = math.inf
    partial = False

    for _ in range(plan.num_probes):
        if tracker.exceeded():
            partial = True
            break
        weights = _renormalize(opt.ask())
        if len(weights) != len(plan.datasets):
            raise ValueError(
                f"optimizer returned {len(weights)} weights; "
                f"expected {len(plan.datasets)}"
            )
        t0 = tracker.elapsed
        try:
            loss = proxy_run(weights)
        except (KeyboardInterrupt, SystemExit):
            raise
        except Exception:
            # Proxy failures are isolated per-candidate (matches v0.33.0
            # #47 CrossDocCollator + v0.40.3 judge_filter_pairs policy).
            # The candidate is recorded with a sentinel high loss so the
            # optimiser sees a valid observation and the run continues.
            import logging
            logging.getLogger(__name__).debug(
                "proxy_run raised for candidate %s", weights, exc_info=True
            )
            opt.tell(weights, _MAX_LOSS)
            continue
        if isinstance(loss, bool) or not isinstance(loss, (int, float)):
            raise TypeError(
                f"proxy_run must return float, got {type(loss).__name__}"
            )
        loss_f = float(loss)
        if not math.isfinite(loss_f):
            # Skip — invalid observation should not poison best-of search.
            opt.tell(weights, _MAX_LOSS)
            continue
        opt.tell(weights, loss_f)
        cand = MixCandidate(
            weights=weights,
            eval_loss=loss_f,
            wall_clock_seconds=tracker.elapsed - t0,
        )
        candidates.append(cand)
        if loss_f < best_loss:
            best_loss = loss_f
            best_weights = weights

    if best_weights is None:
        # No valid observation — pick uniform as graceful fallback.
        n = len(plan.datasets)
        best_weights = tuple([1.0 / n] * n)
        best_loss = math.inf if not candidates else best_loss

    return MixOptimizationReport(
        datasets=plan.datasets,
        candidates=tuple(candidates),
        best_weights=best_weights,
        best_eval_loss=best_loss if math.isfinite(best_loss) else float("nan"),
        partial=partial,
        elapsed_seconds=tracker.elapsed,
    )


# --- Recipe writer --------------------------------------------------------


def render_mix_recipe_yaml(report: MixOptimizationReport) -> str:
    """Render an applied-mixture recipe snippet for human review.

    Produces a YAML fragment suitable for splicing into ``soup.yaml`` under
    ``data:``. Defends against YAML key injection by rejecting newlines and
    null bytes in dataset paths (mirrors v0.46.0 Part A
    ``render_recipe_yaml``).
    """
    if not isinstance(report, MixOptimizationReport):
        raise TypeError(
            "report must be MixOptimizationReport, "
            f"got {type(report).__name__}"
        )
    for path in report.datasets:
        if not isinstance(path, str) or "\n" in path or "\x00" in path:
            raise ValueError(
                "dataset path contains control characters — refusing to "
                "render YAML."
            )
        if len(path) > _MAX_PATH_LEN:
            raise ValueError(
                f"dataset path length {len(path)} exceeds {_MAX_PATH_LEN}"
            )
    lines = ["# Generated by `soup data mix --optimize` (v0.48.0 — BETA)"]
    lines.append(f"# Probes evaluated: {len(report.candidates)}")
    lines.append(
        f"# Best eval loss: {report.best_eval_loss:.6f}"
        if math.isfinite(report.best_eval_loss)
        else "# Best eval loss: (no valid observation)"
    )
    if report.partial:
        lines.append("# Budget exceeded — partial results.")
    lines.append("data:")
    lines.append("  interleave:")
    lines.append("    strategy: probs")
    lines.append("    probs:")
    for w in report.best_weights:
        lines.append(f"      - {w:.6f}")
    lines.append("  train:")
    for path in report.datasets:
        lines.append(f"    - {json.dumps(path)}")
    return "\n".join(lines) + "\n"


def write_mix_recipe(
    report: MixOptimizationReport,
    output_path: str,
    *,
    overwrite: bool = False,
) -> str:
    """Atomically write the rendered recipe to ``output_path``.

    Containment + TOCTOU symlink rejection mirrors v0.46.0 Part A
    ``write_recipe`` and v0.47.0 Part A ``write_forge_dataset``.
    """
    from soup_cli.utils.paths import is_under_cwd

    _check_str_path("output_path", output_path)
    # Symlink check on the ORIGINAL path BEFORE realpath.
    try:
        st = os.lstat(output_path)
    except FileNotFoundError:
        st = None
    except OSError as exc:
        raise ValueError(
            f"output_path is not stat-able: "
            f"{os.path.basename(output_path)!r}"
        ) from exc
    if st is not None and stat.S_ISLNK(st.st_mode):
        raise ValueError(
            f"output_path is a symlink (rejected for safety): "
            f"{os.path.basename(output_path)!r}"
        )
    real = os.path.realpath(output_path)
    if not is_under_cwd(real):
        raise ValueError(
            f"output_path is outside cwd: {os.path.basename(real)!r}"
        )
    if st is not None and not overwrite:
        raise ValueError(
            f"output_path already exists (use overwrite=True): "
            f"{os.path.basename(real)!r}"
        )

    text = render_mix_recipe_yaml(report)
    if len(text.encode("utf-8")) > _MAX_RECIPE_BYTES:
        raise ValueError(
            f"rendered recipe exceeds {_MAX_RECIPE_BYTES} bytes cap"
        )
    parent = os.path.dirname(real) or "."
    os.makedirs(parent, exist_ok=True)
    fd, tmp_path = tempfile.mkstemp(prefix=".mix_recipe.", dir=parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as fh:
            fh.write(text)
        os.replace(tmp_path, real)
    except Exception:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass
        raise
    return real


def load_mix_recipe(path: str) -> Mapping[str, object]:
    """Load + validate a previously-written mix recipe.

    Used by ``soup data mix --apply <recipe.yaml>`` to splice the recommended
    mixture into a target ``soup.yaml``.
    """
    from soup_cli.utils.paths import is_under_cwd

    _check_str_path("path", path)
    # Symlink check on the ORIGINAL path BEFORE realpath.
    try:
        st = os.lstat(path)
    except FileNotFoundError:
        st = None
    except OSError as exc:
        raise ValueError(
            f"recipe path is not stat-able: {os.path.basename(path)!r}"
        ) from exc
    if st is not None and stat.S_ISLNK(st.st_mode):
        raise ValueError(
            f"recipe path is a symlink (rejected for safety): "
            f"{os.path.basename(path)!r}"
        )
    real = os.path.realpath(path)
    if not is_under_cwd(real):
        raise ValueError(
            f"recipe path is outside cwd: {os.path.basename(real)!r}"
        )
    if not os.path.isfile(real):
        raise FileNotFoundError(
            f"recipe not found: {os.path.basename(real)!r}"
        )
    size = os.path.getsize(real)
    if size > _MAX_RECIPE_BYTES:
        raise ValueError(
            f"recipe exceeds {_MAX_RECIPE_BYTES} bytes cap (got {size})"
        )
    import yaml

    with open(real, "r", encoding="utf-8") as fh:
        data = yaml.safe_load(fh)
    if not isinstance(data, Mapping):
        raise ValueError("recipe must be a YAML mapping at top level")
    data_block = data.get("data")
    if not isinstance(data_block, Mapping):
        raise ValueError("recipe missing required 'data:' mapping")
    return data_block