"""LLM-based concept naming using local transformers."""

import json
from typing import Dict, List, Any, Optional

# Global cache for loaded model
_model_cache = {}


def format_naming_prompt(
    prompts: List[str],
    positives: List[str],
    negatives: List[str],
    n_samples: int = 5,
) -> str:
    """Format a prompt for the LLM to name this concept cluster."""

    samples = []
    for i in range(min(n_samples, len(prompts))):
        samples.append(f"""Example {i+1}:
Question: {prompts[i]}
Truthful answer: {positives[i]}
False answer: {negatives[i]}""")

    samples_text = "\n\n".join(samples)

    return f"""You are analyzing a cluster of question-answer pairs from a truthfulness dataset.

These pairs all belong to the same "concept" - they share some common pattern in how truthful vs false answers differ.

Here are sample pairs from this cluster:

{samples_text}

Based on these examples, provide:
1. A short name (2-4 words, snake_case) that captures what makes this cluster distinctive
2. A one-sentence description of what pattern these pairs share

Respond in this exact JSON format:
{{"name": "your_name_here", "description": "Your description here"}}

Only output the JSON, nothing else."""


def get_model_and_tokenizer(model_name: str):
    """Load model and tokenizer, caching for reuse."""
    if model_name in _model_cache:
        return _model_cache[model_name]

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    # Determine device
    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"

    print(f"  Loading {model_name} on {device}...", flush=True)

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16 if device != "cpu" else torch.float32,
        device_map=device,
        trust_remote_code=True,
    )

    _model_cache[model_name] = (model, tokenizer, device)
    return model, tokenizer, device


def call_local_llm(prompt: str, model_name: str = "meta-llama/Llama-3.2-1B-Instruct") -> str:
    """Generate text using local transformers model."""
    import torch

    model, tokenizer, device = get_model_and_tokenizer(model_name)

    # Format as chat message
    messages = [{"role": "user", "content": prompt}]

    # Apply chat template if available
    if hasattr(tokenizer, "apply_chat_template"):
        input_text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
    else:
        input_text = prompt

    inputs = tokenizer(input_text, return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.3,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )

    # Decode only the new tokens
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    return response.strip()


def parse_llm_response(response: str) -> Dict[str, str]:
    """Parse the JSON response from the LLM."""
    # Try to find JSON in the response
    response = response.strip()

    # Look for JSON pattern
    start = response.find("{")
    end = response.rfind("}") + 1

    if start >= 0 and end > start:
        json_str = response[start:end]
        try:
            parsed = json.loads(json_str)
            return {
                "name": parsed.get("name", "unknown_concept"),
                "description": parsed.get("description", "No description available"),
            }
        except json.JSONDecodeError:
            pass

    # Fallback: try to extract name and description manually
    return {
        "name": "unknown_concept",
        "description": response[:200] if response else "No description available",
    }


def name_concept_with_llm(
    prompts: List[str],
    positives: List[str],
    negatives: List[str],
    model: str = "meta-llama/Llama-3.2-1B-Instruct",
) -> Dict[str, str]:
    """Use LLM to generate a meaningful name for a concept cluster."""

    if not prompts:
        return {"name": "empty_cluster", "description": "No samples available"}

    prompt = format_naming_prompt(prompts, positives, negatives)

    try:
        response = call_local_llm(prompt, model)
        return parse_llm_response(response)
    except Exception as e:
        return {
            "name": "naming_failed",
            "description": f"LLM naming failed: {str(e)}",
        }


def name_all_concepts_with_llm(
    concepts: List[Dict[str, Any]],
    pair_texts: Dict[int, Dict[str, str]],
    cluster_labels,
    model: str = "meta-llama/Llama-3.2-1B-Instruct",
) -> List[Dict[str, Any]]:
    """Name all concepts using LLM."""
    import numpy as np

    # Group pairs by concept
    pair_ids = sorted(pair_texts.keys())
    concept_pairs = {i: [] for i in range(len(concepts))}

    for idx, pair_id in enumerate(pair_ids):
        if idx < len(cluster_labels):
            cluster_id = cluster_labels[idx]
            if pair_id in pair_texts:
                concept_pairs[cluster_id].append(pair_texts[pair_id])

    # Name each concept
    for i, concept in enumerate(concepts):
        pairs = concept_pairs.get(i, [])

        if pairs:
            prompts = [p.get("prompt", "") for p in pairs]
            positives = [p.get("positive", "") for p in pairs]
            negatives = [p.get("negative", "") for p in pairs]

            print(f"  Naming concept {i+1} ({len(pairs)} pairs)...", flush=True)
            naming = name_concept_with_llm(prompts, positives, negatives, model)

            concept["name"] = naming["name"]
            concept["description"] = naming["description"]
            concept["llm_named"] = True
        else:
            concept["name"] = f"empty_concept_{i+1}"
            concept["description"] = "Empty concept cluster"
            concept["llm_named"] = False

    return concepts
