import re
import json
from pathlib import Path
from typing import Dict, List
import pandas as pd
from huggingface_hub import HfApi, hf_hub_download


# Task to primary metric mapping
TASK_METRICS = {
    "armenian:finer|0": "ner_accuracy",
    "armenian:pioner|0": "ner_accuracy",
    "armenian:pos|0": "ud_pos_regex_acc",
    "armenian:squad|0": "bleu",
    "armenian:belebele|0": "exact_match_mcqa",
    "armenian:hartak|0": "exact_match_mcqa",
    "armenian:include|0": "exact_match_mcqa",
    "armenian:syndarin|0": "exact_match_mcqa",
    "armenian:dream|0": "exact_match_mcqa",
    "armenian:topic-14class|0": "exact_match_mcqa",
    "armenian:scientific|0": "exact_match_mcqa",
    "armenian:sentiment|0": "exact_match_mcqa",
    "armenian:exam_math|0": "armenian_exam_score",
    "armenian:exam_literature|0": "armenian_exam_score",
    "armenian:exam_history|0": "armenian_exam_score",
    "armenian:email|0": "bleu",
    "armenian:short_sentences_translation|0": "bleu",
    "armenian:conversation|0": "bleu",
    "armenian:arak|0": "bleu",
    "armenian:paraphrase|0": "bleu",
    "armenian:ms_marco|0": "bleu",
    "armenian:mmlu_pro|0": "armenian_mmlu_pro_score",
    "armenian:punctuation|0": "punctuation_accuracy",
    "armenian:space_fix|0": "space_accuracy",
}

# Task categories for grouping
TASK_CATEGORIES = {
    "NER": ["armenian:finer|0", "armenian:pioner|0"],
    "POS": ["armenian:pos|0"],
    "Reading Comprehension": [
        "armenian:squad|0",
        "armenian:belebele|0",
        "armenian:dream|0",
        "armenian:hartak|0",
        "armenian:ms_marco|0",
    ],
    "Classification": [
        "armenian:topic-14class|0",
         "armenian:sentiment|0"
        
    ],
    "MCQA": [ "armenian:include|0",
        "armenian:syndarin|0","armenian:scientific|0"]
    ,
    "Generation": [
        "armenian:email|0",
        "armenian:conversation|0",
        "armenian:arak|0",
        "armenian:paraphrase|0",
    ],
    "Translation": ["armenian:short_sentences_translation|0"],
    "Exams": [
        "armenian:exam_math|0",
        "armenian:exam_literature|0",
        "armenian:exam_history|0",
    ],
    "Text Processing": ["armenian:punctuation|0", "armenian:space_fix|0"],
    "MMLU": ["armenian:mmlu_pro|0"],
}

# Short display names for tasks
TASK_DISPLAY_NAMES = {
    "armenian:finer|0": "FiNER",
    "armenian:pioner|0": "PioNER",
    "armenian:pos|0": "POS",
    "armenian:squad|0": "SQuAD",
    "armenian:belebele|0": "Belebele",
    "armenian:hartak|0": "Hartak - Public Services MCQA",
    "armenian:include|0": "INCLUDE",
    "armenian:syndarin|0": "Syndarin",
    "armenian:dream|0": "DREAM",
    "armenian:topic-14class|0": "Topic-14",
    "armenian:scientific|0": "Scientific",
    "armenian:sentiment|0": "Sentiment",
    "armenian:exam_math|0": "Math Exam",
    "armenian:exam_literature|0": "Lit Exam",
    "armenian:exam_history|0": "Hist Exam",
    "armenian:email|0": "Email Summary",
    "armenian:short_sentences_translation|0": "Short Trans",
    "armenian:conversation|0": "Conversation Summary",
    "armenian:arak|0": "Simple QA",
    "armenian:paraphrase|0": "Paraphrase",
    "armenian:ms_marco|0": "MS MARCO",
    "armenian:mmlu_pro|0": "MMLU-Pro",
    "armenian:punctuation|0": "Punctuation",
    "armenian:space_fix|0": "Space Fix",
}


class ModelHandler:
    def __init__(self, local_results_folder: str = "results"):
        print("[DEBUG] ModelHandler.__init__: Creating HfApi()...")
        self.api = HfApi()
        print("[DEBUG] ModelHandler.__init__: HfApi created")
        self._cache = {}  # Simple cache for fetched results
        self.local_results_folder = Path(local_results_folder)
        print("[DEBUG] ModelHandler.__init__: Initialization complete")

    def _parse_lighteval_results(self, results: Dict) -> Dict:
        """Parse lighteval format results into structured format.

        Keeps raw scores for display (0-20 for exams, 0-100 for BLEU),
        but normalizes to 0-1 when averaging.
        """
        parsed = {
            "tasks": {},
            "categories": {},
        }

        task_results = results.get("results", {})

        for task_key, metrics in task_results.items():
            if task_key in ["all", "armenian:_average|0"]:
                continue

            if task_key in TASK_METRICS:
                primary_metric = TASK_METRICS[task_key]
                if primary_metric in metrics:
                    display_name = TASK_DISPLAY_NAMES.get(task_key, task_key)
                    value = metrics[primary_metric]
                    # Keep raw values for display
                    parsed["tasks"][display_name] = value

        # Calculate category averages
        for category, tasks in TASK_CATEGORIES.items():
            scores = []
            for task in tasks:
                display_name = TASK_DISPLAY_NAMES.get(task, task)
                if display_name in parsed["tasks"]:
                    value = parsed["tasks"][display_name]
                    primary_metric = TASK_METRICS.get(task)

                    # Normalize BLEU to 0-1, keep exam scores as 0-20
                    if primary_metric == "bleu":
                        value = value / 100.0

                    scores.append(value)
            if scores:
                parsed["categories"][category] = sum(scores) / len(scores)

        # Calculate overall average as average of category averages
        # Normalize exam categories (0-20) to 0-1 for final average
        if parsed["categories"]:
            category_scores = []
            for category, avg_score in parsed["categories"].items():
                # Check if this category contains only exam tasks
                tasks_in_category = TASK_CATEGORIES.get(category, [])
                all_exam_tasks = all(
                    TASK_METRICS.get(task) == "armenian_exam_score"
                    for task in tasks_in_category
                )

                # Normalize exam categories from 0-20 to 0-1
                if all_exam_tasks:
                    avg_score = avg_score / 20.0

                category_scores.append(avg_score)

            parsed["average"] = sum(category_scores) / len(category_scores)

        return parsed

    def _fetch_results_from_repo(self, repo_id: str) -> Dict:
        """Fetch results.json from a single HuggingFace repository."""
        # Check cache first
        if repo_id in self._cache:
            return self._cache[repo_id]

        try:
            result_path = hf_hub_download(
                repo_id, filename="results.json", cache_dir=".hf_cache"
            )
            with open(result_path) as f:
                results = json.load(f)
            self._cache[repo_id] = results
            return results
        except FileNotFoundError:
            print(f"  No results.json found in {repo_id}")
            return None
        except Exception as e:
            print(f"  Error fetching results from {repo_id}: {e}")
            return None

    def _is_openrouter_only(self, repo_id: str) -> bool:
        """Check if a repository is OpenRouter-only."""
        return repo_id.lower().startswith("openrouter/")

    def _extract_model_name(self, results: Dict) -> str:
        """Extract model name from results, preferring model_config.model_name."""
        # Try to get from model_config first
        try:
            model_name = (
                results.get("config_general", {})
                .get("model_config", {})
                .get("model_name", "")
            )
            if model_name:
                # Remove 'openrouter/' prefix if present and keep only the model part
                if "/" in model_name:
                    model_name = "/".join(model_name.split("/")[-2:])
                return model_name
        except (KeyError, AttributeError, TypeError):
            pass

        return None

    def _load_local_results(self) -> List[tuple]:
        """Load results from local results folder. Returns list of (model_name, results) tuples."""
        local_models = []

        if not self.local_results_folder.exists():
            return local_models

        print(f"\nLoading local results from {self.local_results_folder}...")
        json_files = list(self.local_results_folder.glob("*.json"))

        if not json_files:
            print(f"  No JSON files found in {self.local_results_folder}")
            return local_models

        print(f"  Found {len(json_files)} local files\n")

        for json_file in sorted(json_files):
            try:
                with open(json_file) as f:
                    results = json.load(f)
                if "results" in results:
                    # Extract model name from config
                    model_name = self._extract_model_name(results)
                    if not model_name:
                        # Fallback to filename stem
                        model_name = json_file.stem
                    local_models.append((model_name, results))
                    print(f"  ✓ Loaded {model_name} from {json_file.name}")
                else:
                    print(f"  ✗ Invalid format in {json_file.name}")
            except json.JSONDecodeError:
                print(f"  ✗ Invalid JSON in {json_file.name}")
            except Exception as e:
                print(f"  ✗ Error reading {json_file.name}: {e}")

        return local_models

    def get_llm_benchmark_data(self) -> pd.DataFrame:
        """Fetch LLM benchmark results from HuggingFace and local results folder."""
        data = []

        # Fetch from HuggingFace
        print("[DEBUG] get_llm_benchmark_data: Starting...")
        print("Fetching models with ArmBench-LLM tag from HuggingFace...")
        try:
            print("[DEBUG] get_llm_benchmark_data: Calling api.list_models()...")
            models = self.api.list_models(filter="ArmBench-LLM")
            print("[DEBUG] get_llm_benchmark_data: api.list_models() returned")
            repositories = [model.modelId for model in models]
            print(f"Found {len(repositories)} models\n")
            print(
                f"[DEBUG] get_llm_benchmark_data: Processing {len(repositories)} repositories"
            )

            for i, repo_id in enumerate(repositories, 1):
                # Skip OpenRouter-only models
                if self._is_openrouter_only(repo_id):
                    print(
                        f"[{i}/{len(repositories)}] Skipping OpenRouter-only: {repo_id}"
                    )
                    continue

                print(f"[{i}/{len(repositories)}] Fetching {repo_id}...")
                results = self._fetch_results_from_repo(repo_id)

                if results and "results" in results:
                    parsed = self._parse_lighteval_results(results)
                    row = {"model_name": repo_id}
                    model_size = self._get_model_size(model_name)
                    row.update(parsed.get("categories", {}))
                    row["Size"] = model_size
                    if "average" in parsed:
                        row["Average"] = parsed["average"]
                    if len(row) > 1:
                        data.append(row)
                        print("  ✓ Added to leaderboard")
                    else:
                        print("  ✗ Invalid results format")
                else:
                    print("  ✗ Could not fetch or parse results")

        except Exception as e:
            print(f"Error fetching from HuggingFace: {e}")
            print(f"[DEBUG] get_llm_benchmark_data: Exception during HF fetch: {e}")

        # Load local results
        print("[DEBUG] get_llm_benchmark_data: Loading local results...")
        local_models = self._load_local_results()
        for model_name, results in local_models:
            parsed = self._parse_lighteval_results(results)
            model_size = self._get_model_size(model_name)
            row = {"model_name": model_name}
            row.update(parsed.get("categories", {}))
            row["Size"] = model_size
            if "average" in parsed:
                row["Average"] = parsed["average"]
            if len(row) > 1:
                # Remove if already exists from HuggingFace (local overrides)
                data = [m for m in data if m["model_name"] != model_name]
                data.append(row)
                print(f"  ✓ Added {model_name} to leaderboard\n")

        print(
            f"[DEBUG] get_llm_benchmark_data: Returning DataFrame with {len(data)} rows"
        )
        return pd.DataFrame(data)

    def get_detailed_results(self) -> Dict[str, pd.DataFrame]:
        """Get detailed task-level results for all models (HuggingFace + local)."""
        print("[DEBUG] get_detailed_results: Starting...")
        print("Fetching detailed results from HuggingFace models...")
        detailed_data = []

        # Fetch from HuggingFace
        try:
            print("[DEBUG] get_detailed_results: Calling api.list_models()...")
            models = self.api.list_models(filter="ArmBench-LLM")
            print("[DEBUG] get_detailed_results: api.list_models() returned")
            repositories = [model.modelId for model in models]
            print(f"Found {len(repositories)} models\n")
            print(
                f"[DEBUG] get_detailed_results: Processing {len(repositories)} repositories"
            )

            for i, repo_id in enumerate(repositories, 1):
                # Skip OpenRouter-only models
                if self._is_openrouter_only(repo_id):
                    print(
                        f"[{i}/{len(repositories)}] Skipping OpenRouter-only: {repo_id}"
                    )
                    continue

                print(f"[{i}/{len(repositories)}] Processing {repo_id}...")
                results = self._fetch_results_from_repo(repo_id)

                if results and "results" in results:
                    parsed = self._parse_lighteval_results(results)
                    row = {"model_name": repo_id}
                    row.update(parsed.get("tasks", {}))
                    if len(row) > 1:
                        detailed_data.append(row)
                        print(f"  ✓ Added {len(parsed.get('tasks', {}))} tasks")
                    else:
                        print("  ✗ No valid tasks")
                else:
                    print("  ✗ Could not fetch results")

        except Exception as e:
            print(f"Error fetching detailed results: {e}")
            print(f"[DEBUG] get_detailed_results: Exception during HF fetch: {e}")

        # Load local results
        print("[DEBUG] get_detailed_results: Loading local results...")
        local_models = self._load_local_results()
        for model_name, results in local_models:
            parsed = self._parse_lighteval_results(results)
            row = {"model_name": model_name}
            row.update(parsed.get("tasks", {}))
            if len(row) > 1:
                # Remove if already exists from HuggingFace (local overrides)
                detailed_data = [
                    m for m in detailed_data if m["model_name"] != model_name
                ]
                detailed_data.append(row)
                print(
                    f"  ✓ Added {model_name} with {len(parsed.get('tasks', {}))} tasks\n"
                )

        print(
            f"[DEBUG] get_detailed_results: Returning with {len(detailed_data)} models"
        )
        return {
            "tasks": pd.DataFrame(detailed_data) if detailed_data else pd.DataFrame(),
        }

    def _get_model_size(self, model_id: str) -> str:

        csv_path = "model_sizes.csv"
        csv_file = Path(csv_path)
        size = None

        if csv_file.is_file():
            df = pd.read_csv(csv_file)
        else:
            df = pd.DataFrame(columns=["Model Name", "Size"])

        try:
            if "Model Name" in df.columns and "Size" in df.columns:
                matching = df[df["Model Name"] == model_id]
                if not matching.empty:
                    size = matching["Size"].iloc[0].strip()

        except Exception as e:
            print(f"Warning: Could not read {csv_path} ({e})")

        lower_id = model_id.lower().replace("/", "-")
        patterns = [
            r"(\d+(?:\.\d+)?)[bB](?:[-_]([a-z]?\d+(?:\.\d+)?[bB]?|[a-z]+\d+[bB]?|[0-9]+x[0-9]+[bB]?))?",
        ]
        if size is None:
            for pattern in patterns:
                match = re.search(pattern, lower_id)
                if match:
                    total = match.group(1)
                    suffix = match.group(2)
                    if suffix:
                        suffix = suffix.upper()
                        if (
                            len(suffix) <= 8
                            and not suffix.isdigit()
                            and not suffix in ["INSTRUCT", "IT", "V2", "BASE"]
                        ):
                            size = f"{total}B-{suffix}"
                        else:
                            size = f"{float(total):g}B"
                    else:
                        size = f"{float(total):g}B"
                    break

        if size is None:
            try:
                api = HfApi()
                info = api.model_info(model_id, timeout=10)

                if hasattr(info, "cardData") and info.cardData:
                    card = info.cardData
                    for key in [
                        "parameters",
                        "num_parameters",
                        "total_params",
                        "model_size",
                        "params",
                    ]:
                        if key in card:
                            val = card[key]
                            if isinstance(val, (int, float)) and val > 1000:
                                size = f"{val / 1e-9:g}B"
                                break
                            if isinstance(val, str):
                                m = re.search(r"(\d+(?:\.\d+)?)\s*[bB]", val.lower())
                                if m:
                                    size = f"{float(m.group(1)):g}B"
                                    break

                if (
                    size is None
                    and hasattr(info, "config")
                    and isinstance(info.config, dict)
                ):
                    cfg = info.config
                    if "num_parameters" in cfg and isinstance(
                        cfg["num_parameters"], (int, float)
                    ):
                        size = f"{cfg['num_parameters'] / 1e-9:g}B"

            except Exception:
                pass

        if size is None:
            size = " - "

        try:
            if model_id not in df["Model Name"].values:
                new_row = pd.DataFrame({"Model Name": [model_id], "Size": [size]})
                df = pd.concat([df, new_row], ignore_index=True)
                df.to_csv(csv_file, index=False)
                print(f"Appended {model_id} → {size} to {csv_path}")

        except Exception as e:
            print(f"Warning: Could not append to {csv_path} ({e})")

        return size