"""Performance interpolators for estimating TTFT, ITL, and throughput.

These interpolators use pre-profiling data (CSV or JSON) to estimate
latency and throughput for given input sequence lengths (ISL).

Supported metrics
-----------------
- **TTFT** (Time To First Token), in seconds
- **ITL** (Inter-Token Latency / decode step time), in seconds/token
- **Throughput** (tokens/second)

Data format
-----------
CSV with header row::

    isl,ttft,itl,throughput
    64,0.05,0.01,450.0
    128,0.08,0.011,430.0
    256,0.14,0.013,390.0
    512,0.24,0.017,330.0
    1024,0.42,0.025,250.0

JSON (list of records)::

    [
        {"isl": 64, "ttft": 0.05, "itl": 0.01, "throughput": 450.0},
        ...
    ]

Issues resolved
---------------
- #6  PerformanceInterpolator.from_csv / from_json — load profiling data
- #7  PerformanceInterpolator.predict_ttft — interpolate TTFT from ISL
- #8  PerformanceInterpolator.predict_itl  — interpolate ITL from ISL
- #9  ThroughputInterpolator.from_csv / from_json — load profiling data
- #10 ThroughputInterpolator.predict_throughput — interpolate throughput
- #11 PerformanceInterpolator.reverse_ttft — find max ISL for target TTFT
"""

from __future__ import annotations

import bisect
import csv
import json
import logging
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Union

logger = logging.getLogger(__name__)


class InterpolationMode(str, Enum):
    """Algorithm used to interpolate between data points.

    Attributes:
        LINEAR: Simple linear interpolation between adjacent points.
        NEAREST: Snap to the nearest data point (no interpolation).
    """

    LINEAR = "linear"
    NEAREST = "nearest"


@dataclass(frozen=True)
class ProfileDataPoint:
    """A single profiling measurement.

    Attributes:
        isl: Input sequence length (tokens).
        ttft: Time-to-first-token in seconds.  ``None`` if not measured.
        itl: Inter-token latency (decode step) in seconds.  ``None`` if
            not measured.
        throughput: Tokens per second.  ``None`` if not measured.
    """

    isl: int
    ttft: float | None = None
    itl: float | None = None
    throughput: float | None = None

    def __post_init__(self) -> None:
        if self.isl < 0:
            raise ValueError(f"isl must be >= 0, got {self.isl}")
        for attr, val in [("ttft", self.ttft), ("itl", self.itl)]:
            if val is not None and val < 0:
                raise ValueError(f"{attr} must be >= 0, got {val}")
        if self.throughput is not None and self.throughput < 0:
            raise ValueError(f"throughput must be >= 0, got {self.throughput}")


@dataclass
class ProfileDataset:
    """A collection of :class:`ProfileDataPoint` objects, sorted by ISL.

    Attributes:
        model_name: Optional model identifier.
        hardware_tag: Optional hardware identifier (e.g. ``"A100-80G"``).
        points: Sorted list of data points.
    """

    model_name: str = ""
    hardware_tag: str = ""
    points: list[ProfileDataPoint] = field(default_factory=list)

    def sorted_points(self) -> list[ProfileDataPoint]:
        """Return points sorted by ISL (ascending)."""
        return sorted(self.points, key=lambda p: p.isl)

    @classmethod
    def from_csv(cls, path: Union[str, Path], **kwargs: Any) -> ProfileDataset:
        """Load profile data from a CSV file.

        The CSV must contain an ``isl`` column.  Additional columns
        ``ttft``, ``itl``, and ``throughput`` are optional.

        Args:
            path: Path to the CSV file.
            **kwargs: Extra fields forwarded to the dataset (``model_name``,
                ``hardware_tag``).

        Returns:
            A :class:`ProfileDataset` populated from the file.

        Raises:
            FileNotFoundError: If the file does not exist.
            ValueError: If ``isl`` column is missing or contains invalid data.
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"Profile CSV not found: {path}")

        points: list[ProfileDataPoint] = []
        with path.open(newline="") as fh:
            reader = csv.DictReader(fh)
            for lineno, row in enumerate(reader, start=2):
                if "isl" not in row:
                    raise ValueError(f"CSV row {lineno} is missing required 'isl' column")
                try:
                    points.append(
                        ProfileDataPoint(
                            isl=int(float(row["isl"])),
                            ttft=float(row["ttft"]) if "ttft" in row and row["ttft"] else None,
                            itl=float(row["itl"]) if "itl" in row and row["itl"] else None,
                            throughput=(
                                float(row["throughput"])
                                if "throughput" in row and row["throughput"]
                                else None
                            ),
                        )
                    )
                except (ValueError, KeyError) as exc:
                    raise ValueError(f"Invalid data at CSV row {lineno}: {exc}") from exc

        dataset = cls(
            model_name=kwargs.get("model_name", ""),
            hardware_tag=kwargs.get("hardware_tag", ""),
            points=points,
        )
        logger.debug("Loaded %d profile points from %s", len(points), path)
        return dataset

    @classmethod
    def from_json(cls, path: Union[str, Path], **kwargs: Any) -> ProfileDataset:
        """Load profile data from a JSON file.

        The JSON file must be a list of objects, each with at minimum an
        ``isl`` key.

        Args:
            path: Path to the JSON file.
            **kwargs: Extra fields forwarded to the dataset.

        Returns:
            A :class:`ProfileDataset` populated from the file.

        Raises:
            FileNotFoundError: If the file does not exist.
            ValueError: If the JSON is malformed or missing required fields.
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"Profile JSON not found: {path}")

        raw: Any = json.loads(path.read_text())
        if not isinstance(raw, list):
            raise ValueError("Profile JSON must be a list of records")

        points: list[ProfileDataPoint] = []
        for idx, record in enumerate(raw):
            if "isl" not in record:
                raise ValueError(f"Record {idx} is missing required 'isl' field")
            try:
                points.append(
                    ProfileDataPoint(
                        isl=int(record["isl"]),
                        ttft=float(record["ttft"]) if "ttft" in record else None,
                        itl=float(record["itl"]) if "itl" in record else None,
                        throughput=(
                            float(record["throughput"]) if "throughput" in record else None
                        ),
                    )
                )
            except (TypeError, KeyError) as exc:
                raise ValueError(f"Invalid JSON record {idx}: {exc}") from exc

        dataset = cls(
            model_name=kwargs.get("model_name", raw[0].get("model_name", "") if raw else ""),
            hardware_tag=kwargs.get("hardware_tag", raw[0].get("hardware_tag", "") if raw else ""),
            points=points,
        )
        logger.debug("Loaded %d profile points from %s", len(points), path)
        return dataset


class PerformanceInterpolator:
    """Interpolate TTFT and ITL from profiling data.

    Supports loading from CSV or JSON and querying via linear or
    nearest-neighbor interpolation.

    Args:
        dataset: Profile data to interpolate over.
        mode: Interpolation algorithm.  Defaults to
            :attr:`InterpolationMode.LINEAR`.

    Example::

        interp = PerformanceInterpolator.from_csv("profiles/qwen2-7b.csv")
        ttft = interp.predict_ttft(512)          # → ~0.24 s
        itl  = interp.predict_itl(512)           # → ~0.017 s/token
        max_isl = interp.reverse_ttft(0.3)       # ISL that hits 0.3 s TTFT
    """

    def __init__(
        self,
        dataset: ProfileDataset,
        mode: InterpolationMode = InterpolationMode.LINEAR,
    ) -> None:
        self._dataset = dataset
        self._mode = mode
        self._sorted: list[ProfileDataPoint] = dataset.sorted_points()
        if not self._sorted:
            raise ValueError("ProfileDataset contains no data points")

    # ------------------------------------------------------------------
    # Factory methods — issue #6 / #9
    # ------------------------------------------------------------------

    @classmethod
    def from_csv(
        cls,
        path: Union[str, Path],
        mode: InterpolationMode = InterpolationMode.LINEAR,
        **dataset_kwargs: Any,
    ) -> PerformanceInterpolator:
        """Create an interpolator by loading profiling CSV data.

        Args:
            path: Path to the CSV file.
            mode: Interpolation algorithm.
            **dataset_kwargs: Forwarded to :meth:`ProfileDataset.from_csv`.

        Returns:
            Configured :class:`PerformanceInterpolator`.
        """
        dataset = ProfileDataset.from_csv(path, **dataset_kwargs)
        return cls(dataset, mode=mode)

    @classmethod
    def from_json(
        cls,
        path: Union[str, Path],
        mode: InterpolationMode = InterpolationMode.LINEAR,
        **dataset_kwargs: Any,
    ) -> PerformanceInterpolator:
        """Create an interpolator by loading profiling JSON data.

        Args:
            path: Path to the JSON file.
            mode: Interpolation algorithm.
            **dataset_kwargs: Forwarded to :meth:`ProfileDataset.from_json`.

        Returns:
            Configured :class:`PerformanceInterpolator`.
        """
        dataset = ProfileDataset.from_json(path, **dataset_kwargs)
        return cls(dataset, mode=mode)

    @classmethod
    def from_points(
        cls,
        points: list[tuple[int, float, float | None, float | None]],
        mode: InterpolationMode = InterpolationMode.LINEAR,
    ) -> PerformanceInterpolator:
        """Create an interpolator from a list of ``(isl, ttft, itl, throughput)`` tuples.

        Args:
            points: Each tuple is ``(isl, ttft, itl, throughput)``; use
                ``None`` for unavailable metrics.
            mode: Interpolation algorithm.

        Returns:
            Configured :class:`PerformanceInterpolator`.
        """
        data_points = [
            ProfileDataPoint(isl=p[0], ttft=p[1], itl=p[2], throughput=p[3]) for p in points
        ]
        dataset = ProfileDataset(points=data_points)
        return cls(dataset, mode=mode)

    # ------------------------------------------------------------------
    # TTFT prediction — issue #7
    # ------------------------------------------------------------------

    def predict_ttft(self, input_seq_len: int) -> float:
        """Predict TTFT (Time To First Token) in seconds for a given ISL.

        Extrapolates linearly beyond the data range if ISL is out of bounds.

        Args:
            input_seq_len: Number of input tokens.

        Returns:
            Predicted TTFT in seconds (>= 0).

        Raises:
            ValueError: If ``input_seq_len`` is negative, or if no point
                in the dataset has a ``ttft`` value.
        """
        if input_seq_len < 0:
            raise ValueError(f"input_seq_len must be >= 0, got {input_seq_len}")
        pts = [p for p in self._sorted if p.ttft is not None]
        if not pts:
            raise ValueError("No data points with 'ttft' values in the dataset")
        return self._interpolate(pts, input_seq_len, attr="ttft")

    # ------------------------------------------------------------------
    # ITL prediction — issue #8
    # ------------------------------------------------------------------

    def predict_itl(self, input_seq_len: int) -> float:
        """Predict ITL (Inter-Token Latency) in seconds/token for a given ISL.

        Args:
            input_seq_len: Number of input tokens.

        Returns:
            Predicted ITL in seconds/token (>= 0).

        Raises:
            ValueError: If ``input_seq_len`` is negative, or if no point
                in the dataset has an ``itl`` value.
        """
        if input_seq_len < 0:
            raise ValueError(f"input_seq_len must be >= 0, got {input_seq_len}")
        pts = [p for p in self._sorted if p.itl is not None]
        if not pts:
            raise ValueError("No data points with 'itl' values in the dataset")
        return self._interpolate(pts, input_seq_len, attr="itl")

    # ------------------------------------------------------------------
    # Throughput prediction — issue #10
    # ------------------------------------------------------------------

    def predict_throughput(self, input_seq_len: int) -> float:
        """Predict throughput (tokens/second) for a given ISL.

        Args:
            input_seq_len: Number of input tokens.

        Returns:
            Predicted throughput in tokens/second (>= 0).

        Raises:
            ValueError: If ``input_seq_len`` is negative, or if no point
                in the dataset has a ``throughput`` value.
        """
        if input_seq_len < 0:
            raise ValueError(f"input_seq_len must be >= 0, got {input_seq_len}")
        pts = [p for p in self._sorted if p.throughput is not None]
        if not pts:
            raise ValueError("No data points with 'throughput' values in the dataset")
        return self._interpolate(pts, input_seq_len, attr="throughput")

    # ------------------------------------------------------------------
    # Reverse TTFT — issue #11
    # ------------------------------------------------------------------

    def reverse_ttft(self, target_ttft: float) -> int:
        """Find the maximum ISL that keeps TTFT under the given target.

        This is the **reverse interpolation** of :meth:`predict_ttft`:
        given a TTFT budget, return the maximum ISL that can meet it.

        Args:
            target_ttft: TTFT upper bound in seconds.

        Returns:
            Maximum input sequence length (tokens) that satisfies the
            TTFT budget.  Returns 0 when the smallest measured ISL
            already exceeds the target.

        Raises:
            ValueError: If ``target_ttft`` is negative, or if no TTFT data
                exists in the dataset.
        """
        if target_ttft < 0:
            raise ValueError(f"target_ttft must be >= 0, got {target_ttft}")
        pts = [p for p in self._sorted if p.ttft is not None]
        if not pts:
            raise ValueError("No data points with 'ttft' values in the dataset")

        # Edge cases
        if pts[0].ttft is not None and pts[0].ttft > target_ttft:
            return 0  # Even the smallest ISL exceeds budget
        if pts[-1].ttft is not None and pts[-1].ttft <= target_ttft:
            return pts[-1].isl  # All ISLs are within budget

        # Binary-search the crossing segment
        lo, hi = 0, len(pts) - 1
        while lo < hi - 1:
            mid = (lo + hi) // 2
            mid_ttft = pts[mid].ttft
            assert mid_ttft is not None
            if mid_ttft <= target_ttft:
                lo = mid
            else:
                hi = mid

        # Linear interpolation within the segment
        p_lo, p_hi = pts[lo], pts[hi]
        assert p_lo.ttft is not None and p_hi.ttft is not None
        ttft_range = p_hi.ttft - p_lo.ttft
        isl_range = p_hi.isl - p_lo.isl
        if ttft_range == 0:
            return p_lo.isl  # flat region — return left endpoint

        frac = (target_ttft - p_lo.ttft) / ttft_range
        return int(p_lo.isl + frac * isl_range)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _interpolate(
        self,
        pts: list[ProfileDataPoint],
        isl: int,
        attr: str,
    ) -> float:
        """Interpolate ``attr`` for the given ISL.

        Args:
            pts: Filtered, ISL-sorted points that have a value for ``attr``.
            isl: Target input sequence length.
            attr: Name of the attribute to interpolate (``ttft``, ``itl``,
                or ``throughput``).

        Returns:
            Interpolated value.
        """
        isls = [p.isl for p in pts]

        if self._mode == InterpolationMode.NEAREST:
            idx = min(range(len(isls)), key=lambda i: abs(isls[i] - isl))
            return float(getattr(pts[idx], attr))

        # LINEAR mode
        if isl <= isls[0]:
            # Extrapolate left (clamp to first point for now)
            return float(getattr(pts[0], attr))
        if isl >= isls[-1]:
            # Extrapolate right using the last two points
            if len(pts) >= 2:
                p1, p2 = pts[-2], pts[-1]
                v1, v2 = float(getattr(p1, attr)), float(getattr(p2, attr))
                if p2.isl == p1.isl:
                    return v2
                slope = (v2 - v1) / (p2.isl - p1.isl)
                return max(0.0, v2 + slope * (isl - p2.isl))
            return float(getattr(pts[-1], attr))

        # Interpolate within range
        pos = bisect.bisect_right(isls, isl) - 1
        p_lo, p_hi = pts[pos], pts[pos + 1]
        v_lo, v_hi = float(getattr(p_lo, attr)), float(getattr(p_hi, attr))
        frac = (isl - p_lo.isl) / (p_hi.isl - p_lo.isl)
        return v_lo + frac * (v_hi - v_lo)

    # ------------------------------------------------------------------
    # Repr
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        n = len(self._sorted)
        isl_range = f"{self._sorted[0].isl}–{self._sorted[-1].isl}" if n else "empty"
        return f"PerformanceInterpolator(n={n}, isl_range={isl_range}, mode={self._mode.value!r})"
