import os
import cv2
import numpy as np
import time
import csv
from concurrent.futures import ThreadPoolExecutor, as_completed
import psutil
import queue

def main():
    with open("iam_words_sample/iam_sample.csv", "r", encoding="utf-8") as f:
        data = list(csv.DictReader(f))[:100]
        
    os.environ["CT_TROCR_HANDWRITTEN"] = os.path.abspath("./local_trocr_large")
    from pipeline.trocr_engine import _recognize_single_line_trocr
    
    print("Pre-warming model...")
    img0 = cv2.imread(data[0]["crop_path"])
    _recognize_single_line_trocr(img0, handwritten=True)
    
    print("Starting simulation of 100 bursts...")
    
    process = psutil.Process()
    start_mem = process.memory_info().rss / 1024 / 1024
    
    # We simulate a bounded queue (maxsize 50) as recommended by the hostile audit
    # But threadpoolexecutor doesn't have a maxsize queue by default.
    # We will simulate the drop logic:
    
    max_queue = 50
    q = queue.Queue(maxsize=max_queue)
    
    latencies = []
    dropped = 0
    
    def worker():
        while True:
            try:
                img_path, start_q_time = q.get(block=False)
            except queue.Empty:
                break
                
            t0 = time.time()
            img = cv2.imread(img_path)
            if img is None: img = np.zeros((100, 100, 3), dtype=np.uint8)
            res = _recognize_single_line_trocr(img, handwritten=True)
            latencies.append(time.time() - t0)
            q.task_done()

    # Create thread pool
    pool = ThreadPoolExecutor(max_workers=2)
    
    start_time = time.time()
    
    # Simulate API burst
    for row in data:
        try:
            q.put_nowait((row["crop_path"], time.time()))
        except queue.Full:
            dropped += 1
            
    # Start workers
    futures = [pool.submit(worker) for _ in range(2)]
    
    max_cpu = 0
    max_q_seen = 0
    
    while q.qsize() > 0:
        max_q_seen = max(max_q_seen, q.qsize())
        max_cpu = max(max_cpu, psutil.cpu_percent(interval=0.1))
        
    q.join()
    
    latencies.sort()
    if latencies:
        p50 = latencies[int(len(latencies)*0.5)]
        p95 = latencies[int(len(latencies)*0.95)]
        p99 = latencies[int(len(latencies)*0.99)]
    else:
        p50 = p95 = p99 = 0
    
    print(f"\n--- SHADOW SIMULATION METRICS ---")
    print(f"Total simulated crops : 100")
    print(f"Queue Max Growth      : {max_q_seen} / {max_queue}")
    print(f"Dropped shadow tasks  : {dropped} ({(dropped/100)*100}%)")
    print(f"Max CPU Utilization   : {max_cpu}%")
    print(f"Peak Memory Footprint : {process.memory_info().rss / 1024 / 1024:.1f} MB (Delta: {process.memory_info().rss / 1024 / 1024 - start_mem:.1f} MB)")
    print(f"P50 Latency           : {p50:.2f} s")
    print(f"P95 Latency           : {p95:.2f} s")
    print(f"P99 Latency           : {p99:.2f} s")
    print(f"Total time to drain   : {time.time() - start_time:.2f} s")

if __name__ == "__main__":
    main()