"""
Trial file retention policy for isogrid optimization.

Rules:
  - NEVER delete:  *.png  (density, mesh, rib figures — full history always kept)
  - NEVER delete:  *.json (params, results, rib profiles)
  - HEAVY files:   NX model copies (.prt, .fem, .sim, .afm, .afem)
                   + Nastran outputs (.op2, .f06, .dat, .log)
  - KEEP heavy:    last KEEP_RECENT trials  +  best KEEP_BEST trials (by objective)
  - STRIP heavy:   all other trial folders

Usage in run_optimization.py:
    rm = TrialRetentionManager(ITER_DIR, keep_recent=10, keep_best=5)
    # after each trial:
    rm.register(trial_number, trial_dir, objective=obj, mass_kg=mass, feasible=ok)
    rm.apply()
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

# Extensions considered "heavy" (copied once, stripped when not in keep set)
# NX model copies + Nastran outputs — everything needed to reproduce / re-open a trial
HEAVY_EXTENSIONS = {".prt", ".fem", ".sim", ".afm", ".afem", ".op2", ".f06", ".dat", ".log"}

# Extensions that are NEVER deleted regardless of retention policy
SAFE_EXTENSIONS = {".png", ".json"}


@dataclass
class _TrialRecord:
    number: int
    path: Path
    objective: float = float("inf")
    mass_kg: float = float("inf")
    feasible: bool = False
    has_heavy: bool = True


class TrialRetentionManager:
    """
    Manages heavy-file retention across trial folders.

    After each trial:
      1. Call register() with the trial's outcome
      2. Call apply() to enforce the keep-recent + keep-best policy
    """

    def __init__(
        self,
        iter_dir: Path,
        keep_recent: int = 10,
        keep_best: int = 5,
    ):
        self.iter_dir = iter_dir
        self.keep_recent = keep_recent
        self.keep_best = keep_best
        self._records: dict[int, _TrialRecord] = {}

    def register(
        self,
        trial_number: int,
        trial_dir: Path,
        objective: float,
        mass_kg: float,
        feasible: bool,
    ) -> None:
        """Register a completed trial so the retention policy can track it."""
        # Detect whether any heavy files currently exist
        has_heavy = False
        if trial_dir.exists():
            has_heavy = any(
                f.is_file() and f.suffix in HEAVY_EXTENSIONS
                for f in trial_dir.iterdir()
            )
        self._records[trial_number] = _TrialRecord(
            number=trial_number,
            path=trial_dir,
            objective=objective,
            mass_kg=mass_kg,
            feasible=feasible,
            has_heavy=has_heavy,
        )

    def apply(self) -> list[int]:
        """
        Enforce retention policy.

        Returns list of trial numbers whose heavy files were stripped.
        """
        if not self._records:
            return []

        all_nums = sorted(self._records.keys())

        # Set 1: most recent N trials
        recent_set: set[int] = set(all_nums[-self.keep_recent :])

        # Set 2: best K trials — feasible first, then lowest objective
        sorted_by_quality = sorted(
            self._records.values(),
            key=lambda r: (0 if r.feasible else 1, r.objective),
        )
        best_set: set[int] = {r.number for r in sorted_by_quality[: self.keep_best]}

        keep_set = recent_set | best_set

        stripped: list[int] = []
        for num, record in self._records.items():
            if num not in keep_set and record.has_heavy:
                n_removed = self._strip_heavy(record)
                if n_removed > 0:
                    stripped.append(num)

        return stripped

    def _strip_heavy(self, record: _TrialRecord) -> int:
        """
        Remove heavy files from a trial folder.

        PNGs and JSONs are NEVER touched.
        Returns the number of files removed.
        """
        if not record.path.exists():
            record.has_heavy = False
            return 0

        removed = 0
        for f in list(record.path.iterdir()):
            if f.is_file() and f.suffix in HEAVY_EXTENSIONS:
                f.unlink()
                removed += 1

        record.has_heavy = False
        return removed

    def summary(self) -> dict:
        """Return a brief status summary."""
        all_nums = sorted(self._records.keys())
        recent_set = set(all_nums[-self.keep_recent :])
        sorted_by_quality = sorted(
            self._records.values(),
            key=lambda r: (0 if r.feasible else 1, r.objective),
        )
        best_set = {r.number for r in sorted_by_quality[: self.keep_best]}
        keep_set = recent_set | best_set

        return {
            "total_trials": len(self._records),
            "keep_recent": self.keep_recent,
            "keep_best": self.keep_best,
            "currently_kept": sorted(keep_set),
            "stripped": sorted(
                n for n, r in self._records.items()
                if not r.has_heavy and n not in keep_set
            ),
        }