import json
import os
import threading
import time
from contextlib import nullcontext
from datetime import datetime, timezone
from typing import Dict, List, Optional, Union

from evalscope.utils import get_logger

logger = get_logger()


class ProgressTracker:
    """
    File-backed hierarchical progress tracker for evaluation and perf pipelines.

    Each nesting level of TqdmLogging contributes one 'stage' entry.  Completed
    stages are retained in the JSON so callers can see full history.

    State schema written to ``<work_dir>/progress.json``::

        {
          "status":   "running" | "completed" | "error",
          "pipeline": "eval" | "perf",
          "total_count": 14042,
          "processed_count": 5200,
          "percent": 37.03,
          "stage": {
            "name": "Running", "label": "eval",
            "current": 1, "total": 3, "status": "running",
            "children": [
              {"name": "Evaluating", "label": "mmlu",
               "current": 1000, "total": 1000, "status": "completed",
               "children": []},
              {"name": "Evaluating", "label": "ceval",
               "current": 320,  "total": 1000, "status": "running",
               "children": []}
            ]
          },
          "updated_at": "2026-03-05T10:05:42Z"
        }

    Nesting is encoded by the tree structure; ``depth`` is used only internally.

    ``write_interval`` controls how often incremental progress updates are
    flushed to disk (in seconds).  Structural events (stage enter/exit, status
    change) are always written immediately regardless of the interval.

    ``total_count`` is an optional pre-computed count of total items to process.
    For ``eval`` pipelines it is computed from benchmark metadata (applying
    ``limit`` and ``repeat`` per subset).  For ``perf`` pipelines it is the
    sum of all ``number`` values.  When provided, ``processed_count`` and
    ``percent`` are derived and written to the JSON on every update.
    """

    # Stage-name prefix used to identify the 'work' stage for each pipeline.
    _PROCESSED_STAGE_NAMES: Dict[str, str] = {
        'eval': 'Evaluating',
        'perf': 'Processing',
    }

    def __init__(
        self,
        work_dir: str,
        pipeline: str = '',
        write_interval: float = 1.0,
        total_count: Optional[int] = None,
    ):
        os.makedirs(work_dir, exist_ok=True)
        self._path = os.path.join(work_dir, 'progress.json')
        self._lock = threading.Lock()
        self._stages: List[dict] = []
        # Maps nesting depth -> index in self._stages for currently *active* stages.
        self._active_depths: Dict[int, int] = {}
        self._status = 'running'
        self._pipeline = pipeline
        self._write_interval = write_interval
        self._last_write_time: float = 0.0
        self._total_count: Optional[int] = total_count
        self._write(force=True)

    # ------------------------------------------------------------------
    # Context manager — wraps the full pipeline lifecycle
    # ------------------------------------------------------------------

    def __enter__(self) -> 'ProgressTracker':
        """Attach this tracker to TqdmLogging and return self."""
        from evalscope.utils.tqdm_utils.tqdm_logging import TqdmLogging
        TqdmLogging.set_tracker(self)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        """Detach tracker and finalise status based on whether an exception occurred."""
        from evalscope.utils.tqdm_utils.tqdm_logging import TqdmLogging
        self.set_status('error' if exc_type is not None else 'completed')
        TqdmLogging.set_tracker(None)
        return False  # never suppress exceptions

    # ------------------------------------------------------------------
    # Status API
    # ------------------------------------------------------------------

    def set_status(self, status: str) -> None:
        with self._lock:
            self._status = status
            self._write(force=True)

    def _enter_stage(self, name: str, total: int, label: str) -> int:
        """Push a new stage onto the stack; returns the nesting depth."""
        with self._lock:
            depth = len(self._active_depths)
            idx = len(self._stages)
            self._stages.append({
                'depth': depth,
                'name': name,
                'label': label,
                'current': 0,
                'total': total,
                'status': 'running',
            })
            self._active_depths[depth] = idx
            self._write(force=True)
        return depth

    def _update_stage(
        self,
        depth: int,
        current: int,
        name: Optional[str] = None,
        label: Optional[str] = None,
    ) -> None:
        """Update progress (and optionally name/label) for the active stage at *depth*.

        Disk writes are throttled to ``write_interval`` seconds to avoid I/O
        overhead on hot paths (e.g. perf benchmark inner loops).
        """
        with self._lock:
            if depth in self._active_depths:
                idx = self._active_depths[depth]
                self._stages[idx]['current'] = current
                if name is not None:
                    self._stages[idx]['name'] = name
                if label is not None:
                    self._stages[idx]['label'] = label
                self._write(force=False)

    def _exit_stage(self, depth: int) -> None:
        """Mark the stage at *depth* (and any orphaned child stages) as completed."""
        with self._lock:
            for d in [d for d in self._active_depths if d >= depth]:
                idx = self._active_depths.pop(d)
                self._stages[idx]['status'] = 'completed'
                self._stages[idx]['current'] = self._stages[idx]['total']
            self._write(force=True)

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _get_processed_count(self) -> Optional[int]:
        """Sum the ``current`` values of all work stages for the active pipeline.

        Returns ``None`` when no matching stages are found or when
        ``total_count`` was not provided (so the field can be omitted).
        """
        if self._total_count is None:
            return None
        stage_name = self._PROCESSED_STAGE_NAMES.get(self._pipeline)
        if stage_name is None:
            return None
        return sum(s['current'] for s in self._stages if s['name'] == stage_name)

    def _build_tree(self, stages: List[dict]) -> Optional[dict]:
        """Convert the flat *stages* list (with depth integers) into a nested tree.

        Each node gains a ``children`` list; the ``depth`` key is dropped from
        the serialised output since the nesting itself encodes the hierarchy.
        """
        if not stages:
            return None

        root: Optional[dict] = None
        stack: List[dict] = []  # stack entries: {'_depth': int, 'node': dict}

        for stage in stages:
            node = {
                'name': stage['name'],
                'label': stage['label'],
                'current': stage['current'],
                'total': stage['total'],
                'status': stage['status'],
                'children': [],
            }
            depth = stage['depth']

            # Pop deeper-or-equal entries so the top of the stack is the parent.
            while stack and stack[-1]['_depth'] >= depth:
                stack.pop()

            if stack:
                stack[-1]['node']['children'].append(node)
            else:
                root = node

            stack.append({'_depth': depth, 'node': node})

        return root

    def _write(self, force: bool = True) -> None:
        """Write progress state to disk.

        When ``force=False`` the write is skipped if less than ``write_interval``
        seconds have elapsed since the last write.  The lock must already be held
        by the caller.
        """
        now = time.monotonic()
        if not force and (now - self._last_write_time) < self._write_interval:
            return
        self._last_write_time = now
        processed = self._get_processed_count()
        percent: Optional[float] = None
        if self._total_count is not None and self._total_count > 0 and processed is not None:
            percent = round(processed / self._total_count * 100, 2)
        state = {
            'status': self._status,
            'pipeline': self._pipeline,
            'total_count': self._total_count,
            'processed_count': processed,
            'percent': percent,
            'stage': self._build_tree(self._stages),
            'updated_at': datetime.now().isoformat(),
        }
        logger.debug(f'Processed / Total: {processed} / {self._total_count} | {percent}%')
        tmp = self._path + '.tmp'
        with open(tmp, 'w') as f:
            json.dump(state, f, indent=2, ensure_ascii=False)
        os.replace(tmp, self._path)


def make_tracker(
    enabled: bool,
    work_dir: str,
    pipeline: str = '',
    write_interval: float = 1.0,
    total_count: Optional[int] = None,
) -> Union['ProgressTracker', 'nullcontext']:
    """Return a ``ProgressTracker`` context manager when *enabled*, otherwise a no-op ``nullcontext``.

    Example::

        with make_tracker(task_config.enable_progress_tracker, outputs.outputs_dir, pipeline='eval', total_count=5000) as tracker:
            ...
    """
    if enabled:
        return ProgressTracker(
            work_dir=work_dir, pipeline=pipeline, write_interval=write_interval, total_count=total_count
        )
    return nullcontext()
