#!/usr/bin/env python3
"""
Part 2 of populate_tasks_backup: LLM-based benchmark tagging and README parsing.
Split from populate_tasks_backup.py to meet 300-line limit.
"""

import random
from typing import Dict, Any, List

from wisent.core.utils import preferred_dtype, resolve_default_device, resolve_device
from wisent.core.utils.config_tools.constants import LLAMA_PAD_TOKEN_ID, DISPLAY_TOP_N_TINY, DISPLAY_TOP_N_MINI


def get_benchmark_tags_with_llama(task_name: str, context_max_length: int, readme_content: str = "") -> List[str]:
    """
    Use Llama-3.1B-Instruct to determine appropriate tags for a benchmark.

    Args:
        task_name: Name of the benchmark
        readme_content: README content describing the benchmark

    Returns:
        List of 3 most appropriate tags
    """
    approved_skills = [
        "coding", "mathematics", "long context", "creative writing",
        "general knowledge", "medical", "law", "science", "history",
        "tool use", "multilingual", "reasoning"
    ]
    approved_risks = [
        "harmfulness", "toxicity", "bias", "hallucination", "violence",
        "adversarial robustness", "sycophancy", "deception"
    ]
    print(f"   Using Llama-3.1B-Instruct to determine tags for '{task_name}'...")
    try:
        from transformers import pipeline
        import torch
        print(f"   Loading Llama-3.1-8B-Instruct pipeline...")
        device_kind = resolve_default_device()
        device_obj = resolve_device(device_kind)
        if device_kind == "cuda" and torch.cuda.is_available():
            print("   Using CUDA device")
        elif device_kind == "mps":
            print("   Using MPS device")
        else:
            print("   Using CPU device")
        torch_dtype = preferred_dtype(device_kind)
        device_map = "auto" if device_kind == "cuda" else None
        if device_kind == "cuda":
            pipeline_device = 0
        elif device_kind == "mps":
            pipeline_device = device_obj
        else:
            pipeline_device = -1
        generator = pipeline(
            "text-generation",
            model="meta-llama/Llama-3.1-8B-Instruct",
            torch_dtype=torch_dtype,
            device_map=device_map,
            device=pipeline_device,
            do_sample=True,
            pad_token_id=LLAMA_PAD_TOKEN_ID
        )
        print(f"   Successfully loaded Llama-3.1-8B-Instruct pipeline")
        description = readme_content[:context_max_length] if readme_content else (
            f"A benchmark called '{task_name}' for evaluating language models."
        )
        user_prompt = f"""Analyze the benchmark and determine exactly 3 tags.

Benchmark: {task_name}
Description: {description}

Available tags:
Skills: {', '.join(approved_skills)}
Risks: {', '.join(approved_risks)}

Instructions:
1. Analyze what this benchmark actually tests
2. Choose EXACTLY 3 tags that best describe what is being evaluated
3. Focus on the primary capabilities/risks being measured
4. Output only the 3 tags, one per line, no explanations

Tags:"""
        formatted_prompt = (
            f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
            f"You are an expert in AI evaluation benchmarks analyzing benchmark "
            f"tasks to determine what specific cognitive abilities they test."
            f"<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n"
            f"{user_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
        )
        print("   Analyzing with Llama...")
        response = generator(formatted_prompt)
        full_response = response[0]['generated_text']
        generated_text = full_response.split(
            "<|start_header_id|>assistant<|end_header_id|>"
        )[-1].strip()
        print(f"   LLM Response: {generated_text}")
        all_approved_tags = approved_skills + approved_risks
        lines = [line.strip() for line in generated_text.split('\n') if line.strip()]
        determined_tags = []
        for line in lines[:DISPLAY_TOP_N_MINI]:
            clean_line = line.strip('- *123456789.').strip()
            for tag in all_approved_tags:
                if tag.lower() == clean_line.lower() or clean_line.lower() in tag.lower():
                    if tag not in determined_tags:
                        determined_tags.append(tag)
                        break
        if len(determined_tags) < DISPLAY_TOP_N_TINY:
            secondary_tags = ["reasoning", "general knowledge", "science"]
            for secondary in secondary_tags:
                if secondary not in determined_tags:
                    determined_tags.append(secondary)
                if len(determined_tags) >= DISPLAY_TOP_N_TINY:
                    break
        determined_tags = determined_tags[:DISPLAY_TOP_N_TINY]
        print(f"   Final LLM-determined tags: {determined_tags}")
        return determined_tags
    except Exception as e:
        print(f"   Error using LLM: {e}")
        print(f"   Using basic content analysis instead...")
        if readme_content:
            return _basic_content_tag_analysis(readme_content)
        return ["reasoning", "general knowledge", "science"]


def _basic_content_tag_analysis(readme_content: str) -> List[str]:
    """Simple content-based tag analysis when LLM is unavailable."""
    content_lower = readme_content.lower()
    determined_tags = []
    if any(word in content_lower for word in ["math", "arithmetic", "calculation"]):
        determined_tags.append("mathematics")
    if any(word in content_lower for word in ["code", "programming", "python"]):
        determined_tags.append("coding")
    if any(word in content_lower for word in ["medical", "health", "clinical"]):
        determined_tags.append("medical")
    if any(word in content_lower for word in ["adversarial", "robust", "challenging"]):
        determined_tags.append("adversarial robustness")
    if any(word in content_lower for word in ["bias", "fairness", "stereotype"]):
        determined_tags.append("bias")
    if any(word in content_lower for word in ["truthful", "hallucination", "factual"]):
        determined_tags.append("hallucination")
    if any(word in content_lower for word in ["multilingual", "cross-lingual"]):
        determined_tags.append("multilingual")
    if "reasoning" not in determined_tags:
        determined_tags.append("reasoning")
    if len(determined_tags) < DISPLAY_TOP_N_TINY and "general knowledge" not in determined_tags:
        determined_tags.append("general knowledge")
    if len(determined_tags) < DISPLAY_TOP_N_TINY:
        determined_tags.append("science")
    return determined_tags[:DISPLAY_TOP_N_TINY]


def get_benchmark_groups_from_readme(task_name: str, timeout: int, context_max_length: int) -> Dict[str, Any]:
    """
    Read the README from lm-eval-harness repository to get benchmark groups
    and use LLM for tags.

    Args:
        task_name: Name of the benchmark (e.g., "superglue")

    Returns:
        Dictionary with group names and LLM-determined tags
    """
    import requests
    task_dir_map = {
        'superglue': 'super_glue', 'super_glue': 'super_glue',
        'super-glue': 'super_glue', 'glue': 'glue', 'mmlu': 'mmlu',
        'truthfulqa': 'truthfulqa', 'hellaswag': 'hellaswag',
        'arc': 'arc', 'winogrande': 'winogrande'
    }
    task_dir = task_dir_map.get(task_name.lower(), task_name.lower())
    readme_url = (
        f"https://raw.githubusercontent.com/EleutherAI/"
        f"lm-evaluation-harness/main/lm_eval/tasks/{task_dir}/README.md"
    )
    try:
        print(f"   Fetching README from: {readme_url}")
        response = requests.get(readme_url, timeout=timeout)
        response.raise_for_status()
        readme_content = response.text
        determined_tags = get_benchmark_tags_with_llama(task_name, context_max_length=context_max_length, readme_content=readme_content)
        groups = []
        lines = readme_content.split('\n')
        in_target_section = False
        for line in lines:
            line = line.strip()
            if (line.lower().startswith('groups') or
                    line.lower().startswith('## groups') or
                    line.lower().startswith('tags') or
                    line.lower().startswith('## tags') or
                    line.lower().startswith('#### groups') or
                    line.lower().startswith('#### tags')):
                in_target_section = True
                continue
            if (in_target_section and line.startswith('##') and
                    not any(x in line.lower() for x in ['groups', 'tags'])):
                break
            if in_target_section and line:
                if line.startswith('* `') and '`' in line and ':' in line:
                    start = line.find('`') + 1
                    end = line.find('`', start)
                    if start > 0 and end > start:
                        group_name = line[start:end].strip()
                        if group_name and '-' in group_name:
                            groups.append(group_name)
                elif ':' in line and not line.startswith('#'):
                    group_name = line.split(':')[0].strip()
                    group_name = group_name.lstrip('*- `').rstrip('`')
                    if group_name and '-' in group_name:
                        groups.append(group_name)
        groups = list(set(groups))
        valid_groups = [g for g in groups if g and len(g) > 3]
        print(f"   Extracted {len(valid_groups)} groups from README: {valid_groups}")
        print(f"   Determined tags from README content: {determined_tags}")
        return {'groups': valid_groups, 'tags': determined_tags}
    except Exception as e:
        print(f"   Failed to fetch README: {e}")
        return {'groups': [], 'tags': []}
