"""Progress estimation mixin for AceStepHandler."""

import json
import os
import threading
import time
from typing import Optional


class ProgressMixin:
    def _get_project_root(self) -> str:
        """Get project root directory path."""
        current_file = os.path.abspath(__file__)
        return os.path.dirname(
            os.path.dirname(
                os.path.dirname(
                    os.path.dirname(
                        os.path.dirname(current_file)
                    )
                )
            )
        )

    def _load_progress_estimates(self) -> None:
        """Load persisted diffusion progress estimates if available."""
        try:
            if os.path.exists(self._progress_estimates_path):
                with open(self._progress_estimates_path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    if isinstance(data, dict) and isinstance(data.get("records"), list):
                        self._progress_estimates = data
        except Exception:
            # Ignore corrupted cache; it will be overwritten on next save.
            self._progress_estimates = {"records": []}

    def _save_progress_estimates(self) -> None:
        """Persist diffusion progress estimates."""
        try:
            os.makedirs(os.path.dirname(self._progress_estimates_path), exist_ok=True)
            with open(self._progress_estimates_path, "w", encoding="utf-8") as f:
                json.dump(self._progress_estimates, f)
        except Exception:
            pass

    def _duration_bucket(self, duration_sec: Optional[float]) -> str:
        if duration_sec is None or duration_sec <= 0:
            return "unknown"
        if duration_sec <= 60:
            return "short"
        if duration_sec <= 180:
            return "medium"
        if duration_sec <= 360:
            return "long"
        return "xlong"

    def _update_progress_estimate(
        self,
        per_step_sec: float,
        infer_steps: int,
        batch_size: int,
        duration_sec: Optional[float],
    ) -> None:
        if per_step_sec <= 0 or infer_steps <= 0:
            return
        record = {
            "device": self.device,
            "infer_steps": int(infer_steps),
            "batch_size": int(batch_size),
            "duration_sec": float(duration_sec) if duration_sec and duration_sec > 0 else None,
            "duration_bucket": self._duration_bucket(duration_sec),
            "per_step_sec": float(per_step_sec),
            "updated_at": time.time(),
        }
        with self._progress_estimates_lock:
            records = self._progress_estimates.get("records", [])
            records.append(record)
            # Keep recent 100 records
            records = records[-100:]
            self._progress_estimates["records"] = records
            self._progress_estimates["updated_at"] = time.time()
            self._save_progress_estimates()

    def _estimate_diffusion_per_step(
        self,
        infer_steps: int,
        batch_size: int,
        duration_sec: Optional[float],
    ) -> Optional[float]:
        # Prefer most recent exact-ish record
        target_bucket = self._duration_bucket(duration_sec)
        with self._progress_estimates_lock:
            records = list(self._progress_estimates.get("records", []))
        if not records:
            return None

        # Filter by device first
        device_records = [r for r in records if r.get("device") == self.device] or records

        # Exact match by steps/batch/bucket
        for r in reversed(device_records):
            if (
                r.get("infer_steps") == infer_steps
                and r.get("batch_size") == batch_size
                and r.get("duration_bucket") == target_bucket
            ):
                return r.get("per_step_sec")

        # Same steps + bucket, scale by batch and duration when possible
        for r in reversed(device_records):
            if r.get("infer_steps") == infer_steps and r.get("duration_bucket") == target_bucket:
                base = r.get("per_step_sec")
                base_batch = r.get("batch_size", batch_size)
                base_dur = r.get("duration_sec")
                if base and base_batch:
                    est = base * (batch_size / base_batch)
                    if duration_sec and base_dur:
                        est *= (duration_sec / base_dur)
                    return est

        # Same steps, scale by batch and duration ratio if available
        for r in reversed(device_records):
            if r.get("infer_steps") == infer_steps:
                base = r.get("per_step_sec")
                base_batch = r.get("batch_size", batch_size)
                base_dur = r.get("duration_sec")
                if base and base_batch:
                    est = base * (batch_size / base_batch)
                    if duration_sec and base_dur:
                        est *= (duration_sec / base_dur)
                    return est

        # Fallback to global median
        per_steps = [r.get("per_step_sec") for r in device_records if r.get("per_step_sec")]
        if per_steps:
            per_steps.sort()
            return per_steps[len(per_steps) // 2]
        return None

    def _start_diffusion_progress_estimator(
        self,
        progress,
        start: float,
        end: float,
        infer_steps: int,
        batch_size: int,
        duration_sec: Optional[float],
        desc: str,
    ):
        """Best-effort progress updates during diffusion using previous step timing."""
        if progress is None or infer_steps <= 0:
            return None, None
        per_step = self._estimate_diffusion_per_step(
            infer_steps=infer_steps,
            batch_size=batch_size,
            duration_sec=duration_sec,
        ) or self._last_diffusion_per_step_sec
        if not per_step or per_step <= 0:
            return None, None
        expected = per_step * infer_steps
        if expected <= 0:
            return None, None
        stop_event = threading.Event()

        def _runner():
            start_time = time.time()
            while not stop_event.is_set():
                elapsed = time.time() - start_time
                frac = min(0.999, elapsed / expected)
                value = start + (end - start) * frac
                try:
                    progress(value, desc=desc)
                except Exception:
                    pass
                stop_event.wait(0.5)

        thread = threading.Thread(target=_runner, name="diffusion-progress", daemon=True)
        thread.start()
        return stop_event, thread