import json
import os
from typing import List, Dict, Any, Tuple
from collections import Counter
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 SemanticQualityAudit:
    def __init__(self):
        # Ground truth for semantic quality
        self.ground_truth = [
            # Colors
            {"query": "maroon", "expected": "red", "acceptable": ["maroon"], "category": "hair_color", "type": "color"},
            {"query": "burgundy", "expected": "red", "acceptable": ["burgundy"], "category": "hair_color", "type": "color"},
            {"query": "cerulean", "expected": "blue", "acceptable": ["ceruledge"], "category": "hair_color", "type": "color"},
            {"query": "azure", "expected": "blue", "acceptable": ["azure_fang"], "category": "hair_color", "type": "color"},
            
            # Scenes
            {"query": "dance floor", "expected": "gala", "acceptable": ["ballroom", "concert stage"], "category": "scene", "type": "synonym"},
            {"query": "royal palace", "expected": "castle", "acceptable": ["temple", "shrine"], "category": "scene", "type": "synonym"},
            
            # Clothing
            {"query": "royal robes", "expected": "evening gown", "acceptable": ["tuxedo", "suit", "cape", "cloak", "crown and royal robes"], "category": "clothing", "type": "related"},
            {"query": "combat gear", "expected": "armor", "acceptable": ["military uniform", "plate armor"], "category": "clothing", "type": "related"},
            
            # Identity
            {"query": "speedy hedgehog", "expected": "Sonic", "category": "character", "type": "identity"},
            {"query": "pink hedgehog", "expected": "Amy Rose", "category": "character", "type": "identity"},
            {"query": "dark hedgehog", "expected": "Shadow", "category": "character", "type": "identity"},
            
            # Typos
            {"query": "crimsn", "expected": "crimson", "acceptable": ["red"], "category": "hair_color", "type": "typo"},
            {"query": "skool uniform", "expected": "school uniform", "category": "clothing", "type": "typo"}
        ]
        
        self.expanded_gt = []
        for item in self.ground_truth:
            self.expanded_gt.append(item)
            if item["type"] == "color":
                self.expanded_gt.append({**item, "query": f"deep {item['query']}"})
                self.expanded_gt.append({**item, "query": f"bright {item['query']}"})
            elif item["type"] == "related":
                self.expanded_gt.append({**item, "query": f"a {item['query']}"})
                self.expanded_gt.append({**item, "query": f"wearing {item['query']}"})
        
        self.adversarial = [
            {"query": "ballroom", "forbidden": "cocktail dress", "category": "scene"},
            {"query": "Amy", "forbidden": "Rouge", "category": "character"},
            {"query": "Shadow", "forbidden": "Sonic", "category": "character"}
        ]

    def evaluate(self, parser: Parser):
        results = {
            "correct": 0,
            "acceptable": 0,
            "incorrect": 0,
            "rejected": 0,
            "total": 0
        }
        failures = []
        confusion_matrix = Counter()

        for item in self.expanded_gt:
            query = item["query"]
            search_results = parser.embedding_engine.search(query, category=item.get("category"), top_k=1)
            
            results["total"] += 1
            if not search_results:
                results["rejected"] += 1
                continue
                
            record, conf = search_results[0]
            if conf < 0.5:
                results["rejected"] += 1
                continue
                
            actual = record.canonical
            actual_cat = record.category
            
            actual_lower = actual.lower().strip()
            expected_lower = item["expected"].lower().strip()
            
            is_correct = (actual_lower == expected_lower) or (expected_lower in actual_lower) or (actual_lower in expected_lower)
            is_acceptable = False
            if not is_correct:
                if "acceptable" in item:
                    is_acceptable = any(acc.lower().strip() in actual_lower for acc in item["acceptable"]) or \
                                    any(actual_lower in acc.lower().strip() for acc in item["acceptable"])
                if not is_acceptable and query.lower().strip() in actual_lower:
                    is_acceptable = True
            
            if item["category"] == "character" and actual_lower != expected_lower:
                is_correct = False
                is_acceptable = False
                is_incorrect = True
            else:
                is_incorrect = not (is_correct or is_acceptable)
            
            if is_correct:
                results["correct"] += 1
            elif is_acceptable:
                results["acceptable"] += 1
            else:
                results["incorrect"] += 1
                failures.append({
                    "query": query,
                    "actual": actual,
                    "actual_category": actual_cat,
                    "expected": item["expected"],
                    "expected_category": item["category"],
                    "confidence": conf
                })
                confusion_matrix[(item["category"], actual_cat)] += 1

        for item in self.adversarial:
            res = parser.embedding_engine.search(item["query"], top_k=5)
            for rec, conf in res:
                if rec.canonical.lower() == item["forbidden"].lower():
                    results["total"] += 1
                    results["incorrect"] += 1
                    failures.append({
                        "query": item["query"],
                        "actual": rec.canonical,
                        "reason": "Identity/Category Swap",
                        "severity": "Critical"
                    })
                    break

        return results, failures, confusion_matrix

if __name__ == "__main__":
    audit = SemanticQualityAudit()
    matcher = ConceptMatcher("data/ontology")
    from src.embeddings.engine import EmbeddingEngine
    engine = EmbeddingEngine(index_dir="data/faiss_indices")
    engine.load_index()
    parser = Parser(matcher, engine)
    
    print("Running semantic quality audit...")
    stats, failures, confusion = audit.evaluate(parser)
    
    total = stats["total"]
    report = {
        "metrics": {
            "correct_retrieval_pct": round(stats["correct"] / total * 100, 2),
            "acceptable_retrieval_pct": round(stats["acceptable"] / total * 100, 2),
            "incorrect_retrieval_pct": round(stats["incorrect"] / total * 100, 2),
            "trustworthy_pct": round((stats["correct"] + stats["acceptable"]) / total * 100, 2)
        },
        "counts": stats
    }
    
    os.makedirs("reports", exist_ok=True)
    with open("reports/semantic_quality.json", "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)
    with open("reports/semantic_failures.json", "w", encoding="utf-8") as f:
        json.dump(failures, f, indent=2)
        
    cm_report = []
    for (expected_cat, actual_cat), count in confusion.items():
        cm_report.append({"expected_category": expected_cat, "actual_category": actual_cat, "count": count})
    with open("reports/semantic_confusion_matrix.json", "w", encoding="utf-8") as f:
        json.dump(cm_report, f, indent=2)
        
    print(f"Audit complete.")
    print(f"Correct + Acceptable: {report['metrics']['trustworthy_pct']}%")
    print(f"Incorrect: {report['metrics']['incorrect_retrieval_pct']}%")