import json
import random
import os
from typing import List, Dict, Any
from src.parser.parser import Parser
from src.ontology.matcher import ConceptMatcher
from src.embeddings.engine import EmbeddingEngine
from src.enrichment.enricher import Enricher

class BenchmarkSuite:
    def __init__(self, ontology_dir: str):
        self.ontology_dir = ontology_dir
        self.concepts = self._load_all_concepts()

    def _load_all_concepts(self) -> Dict[str, List[str]]:
        concepts = {}
        for filename in os.listdir(self.ontology_dir):
            if filename.endswith(".json"):
                path = os.path.join(self.ontology_dir, filename)
                category = filename.replace(".json", "")
                with open(path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    concepts[category] = [item["canonical"] for item in data]
        return concepts

    def generate_benchmarks(self, count: int = 1000) -> List[Dict[str, Any]]:
        benchmarks = []
        
        # Core characters
        sonic_chars = ["Sonic", "Amy Rose", "Shadow", "Rouge", "Blaze", "Silver", "Knuckles", "Tails", "Cream"]
        
        for _ in range(count):
            # Mix strategies
            strategy = random.random()
            
            if strategy < 0.3: # Character focus
                char = random.choice(sonic_chars)
                clothing = random.choice(self.concepts.get("clothing", ["dress"]))
                scene = random.choice(self.concepts.get("scenes", ["beach"]))
                input_text = f"{char} wearing a {clothing} at the {scene}"
                expected = [char, clothing, scene]
            elif strategy < 0.6: # Attribute focus
                style = random.choice(self.concepts.get("styles", ["anime"]))
                hair = random.choice(self.concepts.get("hair_colors", ["blue hair"]))
                eyes = random.choice(self.concepts.get("eye_colors", ["green eyes"]))
                input_text = f"{style} style girl with {hair} and {eyes}"
                expected = [style, hair, eyes]
            else: # Random mix
                parts = []
                expected = []
                for cat in ["clothing", "accessories", "scenes", "emotions", "lighting"]:
                    if random.random() > 0.5 and self.concepts.get(cat):
                        val = random.choice(self.concepts[cat])
                        parts.append(val)
                        expected.append(val)
                input_text = " ".join(parts)
            
            if input_text:
                benchmarks.append({
                    "input": input_text,
                    "expected_entities": expected
                })
                
        return benchmarks

    def evaluate(self, benchmarks: List[Dict[str, Any]], parser: Parser, enricher: Enricher) -> Dict[str, Any]:
        results = []
        total_precision = 0
        total_recall = 0
        
        exact_match_count = 0
        entity_resolved_count = 0
        total_expected_entities = 0
        
        for case in benchmarks:
            ir = parser.parse(case["input"])
            ir = enricher.enrich(ir)
            
            # Extract all found canonicals
            found = set()
            for char in ir.characters:
                found.add(char.name)
                found.update(char.appearance)
                found.update(char.clothing)
                found.update(char.accessories)
                found.update(char.pose)
                found.update(char.expression)
            found.update(ir.scene.locations)
            found.update(ir.scene.lighting)
            found.update(ir.scene.atmosphere)
            found.update(ir.style)
            found.update(ir.effects)
            
            expected = set(case["expected_entities"])
            
            intersection = found.intersection(expected)
            precision = len(intersection) / len(found) if found else 0
            recall = len(intersection) / len(expected) if expected else 1.0
            
            total_precision += precision
            total_recall += recall
            total_expected_entities += len(expected)
            entity_resolved_count += len(intersection)
            
            if recall == 1.0:
                exact_match_count += 1
                
            results.append({
                "input": case["input"],
                "expected": list(expected),
                "found": list(found),
                "precision": precision,
                "recall": recall
            })
            
        avg_precision = total_precision / len(benchmarks)
        avg_recall = total_recall / len(benchmarks)
        f1 = 2 * (avg_precision * avg_recall) / (avg_precision + avg_recall) if (avg_precision + avg_recall) else 0
        
        return {
            "metrics": {
                "precision": round(avg_precision, 4),
                "recall": round(avg_recall, 4),
                "f1": round(f1, 4),
                "exact_match_rate": round(exact_match_count / len(benchmarks), 4),
                "entity_resolution_accuracy": round(entity_resolved_count / total_expected_entities if total_expected_entities else 0, 4)
            },
            "total_cases": len(benchmarks),
            "results": results
        }

if __name__ == "__main__":
    suite = BenchmarkSuite("data/ontology")
    
    print("Generating 1000 benchmarks...")
    benchmarks = suite.generate_benchmarks(1000)
    with open("benchmarks/prompts.json", "w", encoding="utf-8") as f:
        json.dump(benchmarks, f, indent=2)
    
    print("Running evaluation...")
    matcher = ConceptMatcher("data/ontology")
    engine = EmbeddingEngine(index_dir="data/faiss_indices")
    engine.load_index()
    parser = Parser(matcher, engine)
    enricher = Enricher(matcher)
    
    report = suite.evaluate(benchmarks, parser, enricher)
    
    os.makedirs("reports", exist_ok=True)
    with open("reports/benchmark_results.json", "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)
        
    with open("reports/benchmark_results.md", "w", encoding="utf-8") as f:
        f.write("# Benchmark Results\n\n")
        m = report["metrics"]
        f.write(f"- **Total Cases**: {report['total_cases']}\n")
        f.write(f"- **Precision**: {m['precision']}\n")
        f.write(f"- **Recall**: {m['recall']}\n")
        f.write(f"- **F1 Score**: {m['f1']}\n")
        f.write(f"- **Exact Match Rate**: {m['exact_match_rate']}\n")
        f.write(f"- **Entity Resolution Accuracy**: {m['entity_resolution_accuracy']}\n")
    
    print("Benchmark evaluation complete. Reports generated in reports/")