# Copyright 2026 4f4d
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Context enrichment engine — nearest-timestamp matching for race context."""

from __future__ import annotations

import bisect
import logging
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Tuple

from f1radio.models import ClipContext
from f1radio.utils import _parse_iso_datetime

logger = logging.getLogger("f1radio")

# Time windows for nearby event matching
EVENT_WINDOW_SECONDS = 10.0
OVERTAKE_WINDOW_SECONDS = 15.0
PIT_WINDOW_SECONDS = 15.0


def _extract_timestamps(data: List[Dict], date_key: str = "date") -> List[Tuple[float, Dict]]:
    """Extract and sort (timestamp, record) pairs from API data.

    Args:
        data: List of dicts from OpenF1 API.
        date_key: Key containing the datetime string.

    Returns:
        Sorted list of (unix_timestamp, record) tuples.
    """
    result = []
    for record in data:
        date_str = record.get(date_key)
        if not date_str:
            continue
        try:
            dt = _parse_iso_datetime(str(date_str))
            result.append((dt.timestamp(), record))
        except (ValueError, TypeError):
            continue
    result.sort(key=lambda x: x[0])
    return result


def _find_nearest(
    sorted_timestamps: List[Tuple[float, Dict]],
    target_ts: float,
    max_delta: Optional[float] = None,
) -> Optional[Dict]:
    """Find the record with the nearest timestamp to the target.

    Uses binary search (bisect) for O(log n) lookup.

    Args:
        sorted_timestamps: Sorted list of (timestamp, record) pairs.
        target_ts: Target Unix timestamp.
        max_delta: Maximum allowed time difference in seconds. None = no limit.

    Returns:
        Nearest record dict, or None if no match within max_delta.
    """
    if not sorted_timestamps:
        return None

    timestamps = [t[0] for t in sorted_timestamps]
    idx = bisect.bisect_left(timestamps, target_ts)

    # Check the two candidates (idx-1 and idx)
    best = None
    best_delta = float("inf")

    for candidate_idx in [idx - 1, idx]:
        if 0 <= candidate_idx < len(sorted_timestamps):
            delta = abs(sorted_timestamps[candidate_idx][0] - target_ts)
            if delta < best_delta:
                best_delta = delta
                best = sorted_timestamps[candidate_idx][1]

    if max_delta is not None and best_delta > max_delta:
        return None

    return best


def _find_nearby(
    sorted_timestamps: List[Tuple[float, Dict]],
    target_ts: float,
    window_seconds: float,
) -> List[Dict]:
    """Find all records within a time window of the target.

    Args:
        sorted_timestamps: Sorted list of (timestamp, record) pairs.
        target_ts: Target Unix timestamp.
        window_seconds: Time window in seconds (±).

    Returns:
        List of records within the window.
    """
    if not sorted_timestamps:
        return []

    timestamps = [t[0] for t in sorted_timestamps]
    low = bisect.bisect_left(timestamps, target_ts - window_seconds)
    high = bisect.bisect_right(timestamps, target_ts + window_seconds)

    return [sorted_timestamps[i][1] for i in range(low, high)]


def _find_nearest_for_driver(
    sorted_timestamps: List[Tuple[float, Dict]],
    target_ts: float,
    driver_number: int,
    max_delta: float = 120.0,
) -> Optional[Dict]:
    """Find the nearest record for a specific driver.

    Args:
        sorted_timestamps: Sorted list of (timestamp, record) pairs.
        target_ts: Target Unix timestamp.
        driver_number: Driver number to filter by.
        max_delta: Maximum time difference in seconds.

    Returns:
        Nearest matching record, or None.
    """
    driver_data = [
        (ts, rec) for ts, rec in sorted_timestamps
        if rec.get("driver_number") == driver_number
    ]
    return _find_nearest(driver_data, target_ts, max_delta)


class ContextEnricher:
    """Enriches team radio clips with contextual race data.

    Takes pre-fetched context data from multiple OpenF1 endpoints and
    uses nearest-timestamp matching to find what was happening at the
    exact moment of each radio clip.
    """

    def __init__(
        self,
        position_data: List[Dict],
        interval_data: List[Dict],
        stint_data: List[Dict],
        lap_data: List[Dict],
        race_control_data: List[Dict],
        overtake_data: List[Dict],
        pit_data: List[Dict],
        drivers_data: List[Dict],
    ):
        """Initialize with pre-fetched context data.

        All data should be raw lists of dicts from the OpenF1 API.
        """
        self._position = _extract_timestamps(position_data)
        self._intervals = _extract_timestamps(interval_data)
        self._race_control = _extract_timestamps(race_control_data)
        self._overtakes = _extract_timestamps(overtake_data)
        self._pit = _extract_timestamps(pit_data)
        self._laps = _extract_timestamps(lap_data)
        self._drivers = drivers_data

        # Stints don't have per-record timestamps — index by driver + stint
        self._stint_index = self._build_stint_index(stint_data)

        # Pre-compute best lap times per driver
        self._best_laps = self._build_best_laps(lap_data)

    def _build_stint_index(
        self, stint_data: List[Dict]
    ) -> Dict[int, List[Dict]]:
        """Build a driver → stints index.

        Returns dict mapping driver_number → sorted list of stint dicts.
        """
        index: Dict[int, List[Dict]] = {}
        for stint in stint_data:
            dn = stint.get("driver_number")
            if dn is not None:
                index.setdefault(dn, []).append(stint)
        # Sort by stint_number
        for dn in index:
            index[dn].sort(key=lambda s: s.get("stint_number", 0))
        return index

    def _build_best_laps(self, lap_data: List[Dict]) -> Dict[int, float]:
        """Build dict of driver_number → best lap duration (seconds)."""
        best: Dict[int, float] = {}
        for lap in lap_data:
            dn = lap.get("driver_number")
            duration = lap.get("lap_duration")
            if dn is not None and duration is not None:
                try:
                    dur = float(duration)
                    if dn not in best or dur < best[dn]:
                        best[dn] = dur
                except (ValueError, TypeError):
                    pass
        return best

    def _get_stint_at_lap(
        self, driver_number: int, lap_number: Optional[int]
    ) -> Optional[Dict]:
        """Find the stint for a driver at a given lap."""
        stints = self._stint_index.get(driver_number, [])
        if not stints or lap_number is None:
            return stints[-1] if stints else None

        for stint in reversed(stints):
            lap_start = stint.get("lap_start", 0)
            lap_end = stint.get("lap_end")
            if lap_end is not None:
                if lap_start <= lap_number <= lap_end:
                    return stint
            else:
                # Last stint — no end lap
                if lap_number >= lap_start:
                    return stint
        return stints[-1] if stints else None

    def _format_lap_time(self, seconds: Optional[float]) -> Optional[str]:
        """Format lap time in seconds to M:SS.mmm format."""
        if seconds is None:
            return None
        try:
            s = float(seconds)
            minutes = int(s) // 60
            remainder = s - (minutes * 60)
            return f"{minutes}:{remainder:06.3f}"
        except (ValueError, TypeError):
            return None

    def _format_gap(self, gap_value: Any) -> Optional[str]:
        """Format a gap/interval value to a readable string."""
        if gap_value is None:
            return None
        if isinstance(gap_value, (int, float)):
            if gap_value == 0:
                return "Leader"
            return f"+{gap_value:.3f}s"
        return str(gap_value)

    def _format_event(self, event: Dict) -> str:
        """Format a race control event to a readable string."""
        parts = []
        flag = event.get("flag", "")
        if flag:
            parts.append(flag)
        category = event.get("category", "")
        if category and category not in parts:
            parts.append(category)
        message = event.get("message", "")
        if message:
            parts.append(message)
        scope = event.get("scope", "")
        sector = event.get("sector")
        if scope == "Sector" and sector:
            parts.append(f"Sector {sector}")
        return " — ".join(parts) if parts else "Unknown event"

    def _format_overtake(self, overtake: Dict) -> str:
        """Format an overtake event to a readable string."""
        overtaking = overtake.get("overtaking_driver_number", "?")
        overtaken = overtake.get("overtaken_driver_number", "?")
        position = overtake.get("position")

        # Try to resolve driver names
        ot_code = self._resolve_driver_code(overtaking)
        od_code = self._resolve_driver_code(overtaken)

        result = f"{ot_code} overtook {od_code}"
        if position:
            result += f" for P{position}"
        return result

    def _resolve_driver_code(self, driver_number: Any) -> str:
        """Resolve a driver number to a 3-letter code."""
        if isinstance(driver_number, str):
            return driver_number
        for d in self._drivers:
            if d.get("driver_number") == driver_number:
                return d.get("name_acronym", str(driver_number))
        return str(driver_number)

    def _format_pit(self, pit: Dict) -> str:
        """Format a pit stop event to a readable string."""
        dn = pit.get("driver_number", "?")
        code = self._resolve_driver_code(dn)
        duration = pit.get("pit_duration")
        if duration:
            return f"PIT — {code} ({float(duration):.1f}s)"
        return f"PIT — {code}"

    def enrich(self, clip_date: str, driver_number: int) -> ClipContext:
        """Enrich a single clip with contextual race data.

        Args:
            clip_date: ISO 8601 datetime string of the clip.
            driver_number: Driver number.

        Returns:
            ClipContext with all available context data.
        """
        try:
            dt = _parse_iso_datetime(clip_date)
            target_ts = dt.timestamp()
        except (ValueError, TypeError):
            logger.warning(f"Cannot parse clip date: {clip_date}")
            return ClipContext()

        ctx = ClipContext()

        # --- Position ---
        pos_record = _find_nearest_for_driver(
            self._position, target_ts, driver_number, max_delta=120.0
        )
        if pos_record:
            ctx.position = pos_record.get("position")

        # --- Intervals (gap to leader + interval ahead) ---
        int_record = _find_nearest_for_driver(
            self._intervals, target_ts, driver_number, max_delta=120.0
        )
        if int_record:
            ctx.gap_to_leader = self._format_gap(int_record.get("gap_to_leader"))
            ctx.interval_ahead = self._format_gap(int_record.get("interval"))

        # --- Lap data ---
        lap_record = _find_nearest_for_driver(
            self._laps, target_ts, driver_number, max_delta=120.0
        )
        if lap_record:
            ctx.lap_number = lap_record.get("lap_number")
            ctx.last_lap_time = self._format_lap_time(lap_record.get("lap_duration"))

        # --- Best lap ---
        best_dur = self._best_laps.get(driver_number)
        ctx.best_lap_time = self._format_lap_time(best_dur)

        # --- Stint / tyre info ---
        stint = self._get_stint_at_lap(driver_number, ctx.lap_number)
        if stint:
            ctx.compound = stint.get("compound")
            ctx.stint_number = stint.get("stint_number")
            tyre_age_start = stint.get("tyre_age_at_n_laps", 0)
            lap_start = stint.get("lap_start", 0)
            if ctx.lap_number is not None:
                ctx.tyre_age = tyre_age_start + max(0, ctx.lap_number - lap_start)
            else:
                ctx.tyre_age = tyre_age_start

        # --- Nearby race control events ---
        nearby_events = _find_nearby(
            self._race_control, target_ts, EVENT_WINDOW_SECONDS
        )
        ctx.events = [self._format_event(e) for e in nearby_events]

        # --- Nearby overtakes ---
        nearby_overtakes = _find_nearby(
            self._overtakes, target_ts, OVERTAKE_WINDOW_SECONDS
        )
        ctx.overtakes = [self._format_overtake(o) for o in nearby_overtakes]

        # --- Nearby pit stops ---
        nearby_pits = _find_nearby(
            self._pit, target_ts, PIT_WINDOW_SECONDS
        )
        ctx.pit_stops = [self._format_pit(p) for p in nearby_pits]

        return ctx
