import json
import random
import os
import time
import psutil
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 LargeScaleBenchmark:
    def __init__(self, ontology_dir: str):
        self.ontology_dir = ontology_dir
        self.concepts = self._load_all_concepts()

    def _load_all_concepts(self):
        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 = 100000):
        prompts = []
        chars = self.concepts.get("characters", ["Sonic", "Amy Rose", "Shadow"])
        clothing = self.concepts.get("clothing", ["dress", "suit", "armor"])
        scenes = self.concepts.get("scenes", ["beach", "city", "forest"])
        styles = self.concepts.get("styles", ["anime", "realistic"])
        
        # We need a mix of exact matches, aliases, misspellings, and complex phrases
        for i in range(count):
            strategy = random.random()
            char = random.choice(chars) if chars else "Sonic"
            clo = random.choice(clothing) if clothing else "dress"
            sce = random.choice(scenes) if scenes else "beach"
            sty = random.choice(styles) if styles else "anime"
            
            if strategy < 0.2:
                prompts.append(f"{char} in {clo} at {sce}")
            elif strategy < 0.4:
                prompts.append(f"a person wearing luxurious {clo} in a beautiful {sce}")
            elif strategy < 0.6:
                prompts.append(f"Classic {char} wearing elegant {clo}")
            elif strategy < 0.8:
                prompts.append(f"{sty} style character with {clo}")
            else:
                # Add some typos or weird spacing
                prompts.append(f"{char.lower()}   with  {clo.lower()} in  {sce.lower()}")
                
        return prompts

    def run_benchmark(self, prompts: list, parser: Parser, enricher: Enricher):
        latencies = []
        start_mem = psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024
        
        start_time = time.time()
        for text in prompts:
            s = time.time()
            ir = parser.parse(text)
            ir = enricher.enrich(ir)
            latencies.append((time.time() - s) * 1000)
            
        end_time = time.time()
        end_mem = psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024
        
        total_time = end_time - start_time
        throughput = len(prompts) / total_time
        
        latencies.sort()
        p50 = latencies[int(len(latencies) * 0.5)]
        p95 = latencies[int(len(latencies) * 0.95)]
        p99 = latencies[int(len(latencies) * 0.99)]
        
        return {
            "total_prompts": len(prompts),
            "total_time_s": total_time,
            "throughput_req_sec": throughput,
            "latency_ms": {
                "p50": p50,
                "p95": p95,
                "p99": p99,
                "avg": sum(latencies) / len(latencies)
            },
            "memory_mb": {
                "start": start_mem,
                "end": end_mem,
                "diff": end_mem - start_mem
            }
        }

if __name__ == "__main__":
    import argparse
    parser_args = argparse.ArgumentParser()
    parser_args.add_argument("--count", type=int, default=1000, help="Number of prompts to benchmark")
    args = parser_args.parse_args()

    benchmark = LargeScaleBenchmark("data/ontology")
    
    print(f"Generating {args.count} prompts...")
    prompts = benchmark.generate_prompts(args.count)
    
    matcher = ConceptMatcher("data/ontology")
    
    # Try ONNX first, then PyTorch
    try:
        from src.runtime.onnx_runtime import ONNXEmbeddingEngine
        engine = ONNXEmbeddingEngine(index_dir="data/faiss_indices_onnx")
        engine.load_index()
        engine_type = "ONNX"
    except Exception:
        engine = EmbeddingEngine(index_dir="data/faiss_indices")
        engine.load_index()
        engine_type = "PyTorch"
        
    parser = Parser(matcher, engine)
    enricher = Enricher(matcher)
    
    print(f"Running benchmark with {engine_type} engine...")
    report = benchmark.run_benchmark(prompts, parser, enricher)
    
    os.makedirs("reports", exist_ok=True)
    with open("reports/large_scale_benchmark.md", "w", encoding="utf-8") as f:
        f.write("# Large Scale Benchmark\n\n")
        f.write(f"- **Engine**: {engine_type}\n")
        f.write(f"- **Total Prompts**: {report['total_prompts']}\n")
        f.write(f"- **Total Time**: {report['total_time_s']:.2f} s\n")
        f.write(f"- **Throughput**: {report['throughput_req_sec']:.2f} req/s\n\n")
        
        f.write("## Latency\n")
        f.write(f"- **Average**: {report['latency_ms']['avg']:.2f} ms\n")
        f.write(f"- **p50**: {report['latency_ms']['p50']:.2f} ms\n")
        f.write(f"- **p95**: {report['latency_ms']['p95']:.2f} ms\n")
        f.write(f"- **p99**: {report['latency_ms']['p99']:.2f} ms\n\n")
        
        f.write("## Memory\n")
        f.write(f"- **Start**: {report['memory_mb']['start']:.2f} MB\n")
        f.write(f"- **End**: {report['memory_mb']['end']:.2f} MB\n")
        f.write(f"- **Diff**: {report['memory_mb']['diff']:.2f} MB\n")
        
    print(json.dumps(report, indent=2))
    print("Benchmark complete. Report saved to reports/large_scale_benchmark.md")