"""
Simple data loader for OpenHands Index leaderboard.
Loads JSONL files from local directory or GitHub repository.
"""
import os
import pandas as pd
import json
from pathlib import Path


class SimpleLeaderboardViewer:
    """Simple replacement for agent-eval's LeaderboardViewer."""
    
    def __init__(self, data_dir: str, config: str, split: str):
        """
        Args:
            data_dir: Path to data directory
            config: Config name (e.g., "1.0.0-dev1")
            split: Split name (e.g., "validation" or "test")
        """
        self.data_dir = Path(data_dir)
        self.config = config
        self.split = split
        self.config_path = self.data_dir / config
        
        # Load suite configuration
        config_file = self.config_path / "agenteval.json"
        if config_file.exists():
            with open(config_file) as f:
                suite_config = json.load(f)
                self.suite_config = suite_config["suite_config"]
        else:
            self.suite_config = {
                "name": "openhands-index",
                "version": config,
                "splits": []
            }
        
        # Build tag map from config
        self.tag_map = {}
        for split_config in self.suite_config.get("splits", []):
            if split_config["name"] == split:
                for task in split_config.get("tasks", []):
                    for tag in task.get("tags", []):
                        if tag not in self.tag_map:
                            self.tag_map[tag] = []
                        self.tag_map[tag].append(task["name"])
        
    def _load(self):
        """Load the JSONL file for the split and return DataFrame and tag map."""
        jsonl_file = self.config_path / f"{self.split}.jsonl"
        
        if not jsonl_file.exists():
            # Return empty dataframe with error message
            return pd.DataFrame({
                "Message": [f"No data found for split '{self.split}'. Expected file: {jsonl_file}"]
            }), {}
        
        try:
            # Read JSONL file
            records = []
            with open(jsonl_file, 'r') as f:
                for line in f:
                    if line.strip():
                        records.append(json.loads(line))
            
            if not records:
                return pd.DataFrame({
                    "Message": [f"No data in file: {jsonl_file}"]
                }), {}
            
            df = pd.DataFrame(records)
            return df, self.tag_map
        except Exception as e:
            return pd.DataFrame({
                "Message": [f"Error loading data: {e}"]
            }), {}
    
    def get_dataframe(self):
        """Get the raw dataframe."""
        df, _ = self._load()
        return df


def load_mock_data_locally(data_dir: str = "mock_results"):
    """
    Load mock data from local directory for testing.
    
    Args:
        data_dir: Path to mock results directory
        
    Returns:
        Dictionary mapping split names to SimpleLeaderboardViewer instances
    """
    viewers = {}
    data_path = Path(data_dir)
    
    if not data_path.exists():
        print(f"Warning: Mock data directory '{data_dir}' not found")
        return viewers
    
    # Find all config directories
    for config_dir in data_path.iterdir():
        if config_dir.is_dir():
            config_name = config_dir.name
            
            # Find all JSONL files (each represents a split)
            for jsonl_file in config_dir.glob("*.jsonl"):
                split_name = jsonl_file.stem
                viewer = SimpleLeaderboardViewer(
                    data_dir=str(data_path),
                    config=config_name,
                    split=split_name
                )
                viewers[split_name] = viewer
    
    return viewers