agents/data_agent.py

"""
KRONECTOR - DataAgent.

Plain Python agent that turns a natural-language F1 query into a DataFrame
compatible with ml.predict.predict_dataframe().

Usage:
    agent = DataAgent()
    df = agent.query("What was Verstappen's win probability at Monaco 2023?")
    # Returns a DataFrame with one row matching the query
"""

from __future__ import annotations

import json
import os
import re
from pathlib import Path
from typing import Any, Callable, NotRequired, TypedDict

import pandas as pd
from groq import Groq


DEFAULT_DATA_PATH = Path("data_output/fastf1_races.parquet")
# Current available model (as of June 2026)
GROQ_MODEL = "llama-3.3-70b-versatile"


class QueryIntent(TypedDict):
    season: int
    round: NotRequired[int | None]
    grand_prix: NotRequired[str | None]
    driver_id: NotRequired[str | None]
    driver_name: NotRequired[str | None]
    query_intent: NotRequired[str]


class PredictionInputRow(TypedDict, total=False):
    season: int
    round: int
    driver_id: str
    driver_name: str
    team: str
    grid_position: float
    finish_position: float
    circuit_id: str
    sector_1_time: float
    sector_2_time: float
    sector_3_time: float
    avg_lap_time_practice: float
    tire_compound: float
    tire_age_laps: float
    fresh_tire: float
    pit_stop_count: float
    team_pit_speed: float
    weather_temp_track: float
    weather_rainfall: float
    championship_standing: float
    driver_form_last3: float
    safety_car_probability: float
    telemetry_available: bool
    regulation_era: str
    track_type: str
    win_probability: int


class DataAgentOutput(TypedDict):
    query: str
    intent: QueryIntent
    rows: list[PredictionInputRow]
    dataframe: pd.DataFrame


IntentParser = Callable[[str], QueryIntent]


def _extract_json_object(text: str) -> dict[str, Any]:
    """Extract a JSON object from an LLM response."""
    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if not match:
        raise ValueError(f"Groq response did not contain JSON: {text!r}")
    return json.loads(match.group(0))


def _coerce_int(value: Any, field_name: str) -> int:
    if value is None:
        raise ValueError(f"Missing required query field: {field_name}")
    try:
        return int(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"Invalid {field_name}: {value!r}") from exc


def _normalize_driver_id(value: Any) -> str | None:
    if value is None:
        return None
    driver_id = str(value).strip()
    if not driver_id:
        return None
    return driver_id.upper()


def parse_query_with_groq(query: str) -> QueryIntent:
    """
    Use Groq llama3-70b-8192 to extract season, round, and driver intent.

    Expected JSON shape:
        {"season": 2023, "round": 1, "driver_id": "VER", "driver_name": "Max Verstappen"}
    """
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        raise RuntimeError("GROQ_API_KEY is required for DataAgent query parsing")

    client = Groq(api_key=api_key)
    prompt = (
        "Extract F1 prediction intent from the user query. Return only JSON with "
        "keys season (int), round (int, or null if the user asks by race name), "
        "grand_prix (str, e.g. 'Canadian Grand Prix', or null), "
        "driver_id, driver_name. Use a FastF1 three-letter "
        "driver_id when clear, otherwise null. Do not add prose.\n\n"
        f"Query: {query}"
    )
    response = client.chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {
                "role": "system",
                "content": "You are a strict JSON parser for Formula 1 race queries.",
            },
            {"role": "user", "content": prompt},
        ],
        temperature=0,
    )
    content = response.choices[0].message.content or ""
    parsed = _extract_json_object(content)
    
    # Season is required
    season = _coerce_int(parsed.get("season"), "season")
    
    # We need either round or grand_prix
    round_num = parsed.get("round")
    if round_num is not None:
        round_num = _coerce_int(round_num, "round")
        
    return {
        "season": season,
        "round": round_num,
        "grand_prix": parsed.get("grand_prix"),
        "driver_id": _normalize_driver_id(parsed.get("driver_id")),
        "driver_name": parsed.get("driver_name"),
    }


def _load_feature_data(data_path: str | Path) -> pd.DataFrame:
    path = Path(data_path)
    if not path.exists():
        raise FileNotFoundError(f"Feature data not found: {path}")
    return pd.read_parquet(path)


def _filter_by_intent(df: pd.DataFrame, intent: QueryIntent, filter_driver: bool = True) -> pd.DataFrame:
    rows = df[df["season"] == intent["season"]].copy()
    
    round_num = intent.get("round")
    grand_prix = intent.get("grand_prix")
    
    if round_num is not None:
        rows = rows[rows["round"] == round_num]
    elif grand_prix is not None:
        name_lower = str(grand_prix).strip().lower()
        rows = rows[
            rows["circuit_id"].astype(str).str.lower().str.contains(
                re.escape(name_lower), na=False
            )
        ]
        
        # If we matched a grand prix string, but it happened multiple times in the season 
        # (extremely rare, e.g. 2020 Austrian/Styrian), just take the first round that matched.
        if not rows.empty:
            matched_round = rows["round"].iloc[0]
            rows = rows[rows["round"] == matched_round]
    else:
        raise ValueError("Either 'round' or 'grand_prix' must be provided in the query.")

    if rows.empty:
        raise ValueError(
            f"No rows found for season={intent['season']} round={round_num} grand_prix={grand_prix}"
        )

    if not filter_driver:
        return rows.reset_index(drop=True)

    driver_id = _normalize_driver_id(intent.get("driver_id"))
    driver_name = intent.get("driver_name")

    if driver_id:
        rows = rows[rows["driver_id"].astype(str).str.upper() == driver_id]
    elif driver_name:
        name_lower = str(driver_name).strip().lower()
        rows = rows[
            rows["driver_name"].astype(str).str.lower().str.contains(
                re.escape(name_lower), na=False
            )
        ]

    if rows.empty:
        raise ValueError(f"No rows matched driver intent: {intent}")

    return rows.reset_index(drop=True)


def build_prediction_dataframe(intent: QueryIntent, data_path: str | Path) -> pd.DataFrame:
    """Build a predict_dataframe-compatible DataFrame from parsed intent."""
    df = _load_feature_data(data_path)
    return _filter_by_intent(df, intent)


def data_agent(
    query: str,
    data_path: str | Path = DEFAULT_DATA_PATH,
    parser: IntentParser | None = None,
) -> DataAgentOutput:
    """
    Parse a natural-language query and return prediction-ready rows.

    Args:
        query: Natural-language F1 query.
        data_path: Parquet dataset with fastf1_pipeline/feature schema.
        parser: Optional parser override for tests or offline use.
    """
    parse = parser or parse_query_with_groq
    intent = parse(query)
    dataframe = build_prediction_dataframe(intent, data_path)

    return {
        "query": query,
        "intent": intent,
        "rows": dataframe.to_dict(orient="records"),
        "dataframe": dataframe,
    }


def main() -> None:
    """CLI entry point for the data agent."""
    import argparse

    parser = argparse.ArgumentParser(
        description="KRONECTOR DataAgent: Natural language to F1 prediction data"
    )
    parser.add_argument("query", help="Natural language F1 query")
    parser.add_argument(
        "--data-path",
        default=str(DEFAULT_DATA_PATH),
        help=f"Path to feature dataset (default: {DEFAULT_DATA_PATH})",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output rows as JSON instead of DataFrame repr",
    )
    args = parser.parse_args()

    try:
        result = data_agent(args.query, data_path=args.data_path)
        print(f"Query: {result['query']}")
        print(f"Parsed Intent: {json.dumps(result['intent'], indent=2)}")
        print(f"Rows found: {len(result['rows'])}\n")

        if args.json:
            print(json.dumps(result["rows"], indent=2))
        else:
            print(result["dataframe"].to_string(index=False))

    except Exception as e:
        print(f"Error: {e}", file=__import__("sys").stderr)
        raise


if __name__ == "__main__":
    main()