from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Dict, Any, Iterator, Optional

import json
import pandas as pd

from openevalkit.run import Run



@dataclass
class Dataset:
    """
    A collection of runs for evaluation

    provides methods to load from files, filter, split and iterate
    
    """
    runs: List[Run]
    metadata: Dict[str, Any] = field(default_factory=dict)
    version: str = "1.0"

    @classmethod
    def from_runs(cls, runs: List[Run], **metadata) -> Dataset:
        "Creates a dataset from a list of runs"
        return cls(runs=runs, metadata=metadata)
    
    def __len__(self) -> int:
        return len(self.runs)
    
    def __iter__(self) -> Iterator[Run]:
        return iter(self.runs)

    def __getitem__(self, index: int) -> Run:
        return self.runs[index]

    @classmethod
    def from_jsonl(
        cls, 
        path: str,
        input_field: str = "input",
        output_field: str = "output",
        reference_field: Optional[str] = "reference",
        id_field: Optional[str] = None,  # Auto-generate if None
        metadata_fields: List[str] = [],  # Specific fields for metadata
        metrics_fields: List[str] = []    # Specific fields for metrics
    ) -> Dataset:
        """
        Load dataset from JSONL file.

        Args:
            path: Path to JSONL file
            input_field: Field name for input (default: "input")
            output_field: Field name for output (default: "output")
            reference_field: Field name for reference (default: "reference")
            id_field: Field name for ID (default: None = auto-generate)
            metadata_fields: Specific fields to extract as metadata (default: [] = auto-extract)
            metrics_fields: Specific fields to extract as metrics (default: [])
            
        Returns:
            Dataset with runs from file
            
        Example JSONL format:
            {"id": "1", "input": "What is 2+2?", "output": "4", "reference": "4"}
            {"id": "2", "input": "Capital?", "output": "Paris", "user_id": "u123", "latency": 0.5}
        """
        runs = []

        with open(path, 'r', encoding='utf-8') as f:
            for idx, line in enumerate(f):
                if not line.strip():  # Skip empty lines
                    continue
                data = json.loads(line)

                # Get ID: use id_field or auto-generate 
                run_id = str(data[id_field]) if id_field and id_field in data else str(idx)

                # Get reference
                reference = data.get(reference_field) if reference_field else None

                # Extract metadata
                if metadata_fields:
                    # Explicit: only extract specified fields
                    metadata = {k: data[k] for k in metadata_fields if k in data}
                else:
                    # Auto: extract everything except known fields
                    known_fields = {input_field, output_field, reference_field, id_field}
                    metadata = {k: v for k, v in data.items() if k not in known_fields}
                
                # Extract metrics
                metrics = {}
                for field in metrics_fields:
                    if field in data:
                        metrics[field] = float(data[field])

                run = Run(
                    id=run_id, 
                    input=data[input_field], 
                    output=data[output_field], 
                    reference=reference,
                    metadata=metadata,
                    metrics=metrics if metrics else None
                )
                runs.append(run)

        return cls(
            runs=runs,
            metadata={"source": path, "format": "jsonl", "num_runs": len(runs)}
        )

    @classmethod
    def from_csv(
            cls, 
                path: str,
                input_col: str = "input",
                output_col: str = "output",
                reference_col: Optional[str] = "reference",
                id_col: Optional[str] = None,  # Auto-generate if None
                metadata_cols: List[str] = [],
                metrics_cols: List[str] = [],
    ) -> Dataset:
        """
        Load dataset from CSV file.
        
        Args:
            path: Path to CSV file
            input_col: Column name for input
            output_col: Column name for output
            reference_col: Column name for reference (optional)
            id_col: Column name for ID (auto-generate if None)
            metadata_cols: Columns to include in metadata
            metrics_cols: Columns to include in metrics
            
        Returns:
            Dataset with runs from CSV
            
        Example:
            dataset = Dataset.from_csv(
                "data.csv",
                input_col="question",
                output_col="answer",
                reference_col="expected",
                metadata_cols=["user_id"],
                metrics_cols=["latency"]
            )
        """

        df = pd.read_csv(path)
        runs = []

        for idx, row in df.iterrows():
            # Get ID
            run_id = str(row[id_col]) if id_col else str(idx)

            # Get reference (check if column exists and value is not NaN) 
            reference = None
            if reference_col and reference_col in df.columns and pd.notna(row.get(reference_col)):
                reference = row[reference_col]
            
            # Extract metadata
            metadata = {col: row.get(col) for col in metadata_cols if col in df.columns}

            # Extract metrics
            metrics = {}
            for col in metrics_cols:
                if col in df.columns and pd.notna(row[col]):
                    metrics[col] = float(row[col])
            
            run = Run(
                id=run_id,
                input=row[input_col],
                output=row[output_col],
                reference=reference,
                metadata=metadata,
                metrics=metrics if metrics else None
                )
            runs.append(run)

        return cls(
        runs=runs,
        metadata={"source": path, "format": "csv", "num_runs": len(runs)}
        )




        