import json
import random
import os
from collections import Counter
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 ResolutionAudit:
    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_prompts(self, count: int = 10000) -> List[str]:
        prompts = []
        chars = ["Sonic", "Amy Rose", "Shadow", "Rouge", "Blaze", "Silver", "Knuckles", "Tails", "Cream"]
        clothing = self.concepts.get("clothing", ["dress"])
        scenes = self.concepts.get("scenes", ["beach"])
        styles = self.concepts.get("styles", ["anime"])
        
        for _ in range(count):
            strategy = random.random()
            if strategy < 0.3:
                # Exact matches
                prompts.append(f"{random.choice(chars)} in {random.choice(clothing)} at {random.choice(scenes)}")
            elif strategy < 0.6:
                # Paraphrased / Semantic (simulated)
                prompts.append(f"a hedgehog wearing luxurious {random.choice(clothing)} in a beautiful {random.choice(scenes)}")
            elif strategy < 0.8:
                # Modifiers / Compounds
                prompts.append(f"Classic {random.choice(chars)} wearing elegant {random.choice(clothing)}")
            else:
                # Complex / Random
                prompts.append(f"{random.choice(styles)} style girl with colored hair and {random.choice(clothing)}")
        return prompts

    def run_audit(self, prompts: List[str], parser: Parser, enricher: Enricher):
        stats = Counter()
        semantic_concepts = Counter()
        gap_mining = []
        
        total_resolved = 0
        
        for text in prompts:
            ir = parser.parse(text)
            ir = enricher.enrich(ir)
            
            p_stats = ir.trace.get("resolution_stats", {})
            for method, count in p_stats.items():
                stats[method] += count
                total_resolved += count
                
            # Mining gaps and semantic dependency
            for term, res in ir.trace.get("resolved", {}).items():
                if res["method"] == "semantic":
                    semantic_concepts[(res["canonical"], res.get("category", "unknown"))] += 1
                    gap_mining.append({
                        "term": term,
                        "canonical": res["canonical"],
                        "confidence": res["confidence"]
                    })
        
        # Distribution
        distribution = {method: round(count / total_resolved * 100, 2) for method, count in stats.items()}
        
        return {
            "distribution": distribution,
            "semantic_dependency": semantic_concepts.most_common(50),
            "gap_mining": gap_mining[:100]
        }

if __name__ == "__main__":
    audit = ResolutionAudit("data/ontology")
    print("Generating 1000 prompts for audit...")
    prompts = audit.generate_prompts(1000)
    
    matcher = ConceptMatcher("data/ontology")
    engine = EmbeddingEngine(index_dir="data/faiss_indices")
    engine.load_index()
    parser = Parser(matcher, engine)
    enricher = Enricher(matcher)
    
    print("Running full system audit...")
    results = audit.run_audit(prompts, parser, enricher)
    
    print("Running ablation study (Embeddings Disabled)...")
    parser_no_emb = Parser(matcher, embedding_engine=None)
    
    import time
    
    # Measure Full
    start = time.time()
    audit.run_audit(prompts, parser, enricher)
    full_time = time.time() - start
    
    # Measure Ablation
    start = time.time()
    results_no_emb = audit.run_audit(prompts, parser_no_emb, enricher)
    no_emb_time = time.time() - start
    
    os.makedirs("reports", exist_ok=True)
    
    # Task 5 report
    with open("reports/embedding_ablation.md", "w", encoding="utf-8") as f:
        f.write("# Embedding Ablation Study\n\n")
        f.write("| Metric | Full System | No Embeddings |\n")
        f.write("| :--- | :--- | :--- |\n")
        f.write(f"| Avg Latency | {full_time/1000*1000:.2f}ms | {no_emb_time/1000*1000:.2f}ms |\n")
        # We use 'exact_canonical' + 'alias' as a proxy for recall stability
        full_res = results["distribution"].get("exact_canonical", 0) + results["distribution"].get("alias", 0)
        no_emb_res = results_no_emb["distribution"].get("exact_canonical", 0) + results_no_emb["distribution"].get("alias", 0)
        f.write(f"| Deterministic Resolution % | {full_res}% | {no_emb_res}% |\n")
        f.write(f"| Semantic Fallback % | {results['distribution'].get('semantic', 0)}% | 0% |\n")
    
    # Task 2 report
    with open("reports/resolution_source_distribution.json", "w", encoding="utf-8") as f:
        json.dump(results["distribution"], f, indent=2)
        
    # Task 3 report
    with open("reports/semantic_dependency.md", "w", encoding="utf-8") as f:
        f.write("# Semantic Dependency Report\n\n")
        f.write("| Concept | Category | Frequency |\n")
        f.write("| :--- | :--- | :--- |\n")
        for (concept, cat), freq in results["semantic_dependency"]:
            f.write(f"| {concept} | {cat} | {freq} |\n")
            
    # Task 4 report
    with open("reports/ontology_gap_mining.md", "w", encoding="utf-8") as f:
        f.write("# Ontology Gap Mining Report\n\n")
        f.write("| Raw Term | Resolved Canonical | Confidence |\n")
        f.write("| :--- | :--- | :--- |\n")
        for item in results["gap_mining"]:
            f.write(f"| {item['term']} | {item['canonical']} | {item['confidence']} |\n")
            
    print("Audit complete. Reports generated in reports/")