"""Load test: Simulate 100 students submitting homework simultaneously."""

import os
import time
import shutil
import tempfile
import zipfile
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from threading import Lock
from dotenv import load_dotenv
from ADFMentor import ADFMentor

load_dotenv()

# Thread-safe counters
results_lock = Lock()
success_count = 0
fail_count = 0
total_score = 0


def create_mentor_or_report() -> ADFMentor | None:
    """Create mentor safely; return None and print actionable message if API key is missing."""
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        print("❌ GEMINI_API_KEY is missing.")
        print("   Add GEMINI_API_KEY to .env and rerun.")
        return None
    try:
        return ADFMentor(api_key=api_key)
    except Exception as e:
        print(f"❌ Failed to initialize ADFMentor: {e}")
        return None


def summarize_results(results, label: str, elapsed: float):
    """Print consistent summary metrics for a test run."""
    success_results = [r for r in results if r.get("status") == "success"]
    failed_results = [r for r in results if r.get("status") == "failed"]
    total = len(results)

    scores = [r.get("score", 0) for r in success_results]
    latencies = [r.get("time", 0) for r in success_results]

    print(f"\n\n✅ {label} Complete!")
    print(f"   Time taken: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)")
    print(f"   Success: {len(success_results)}/{total}")
    print(f"   Failed: {len(failed_results)}/{total}")

    if scores:
        print(f"   Average score: {sum(scores)/len(scores):.1f}/100")
        print(f"   Min/Max score: {min(scores)}/{max(scores)}")

    if latencies:
        sorted_latencies = sorted(latencies)
        p95_index = max(0, min(len(sorted_latencies) - 1, int(0.95 * len(sorted_latencies)) - 1))
        p95_latency = sorted_latencies[p95_index]
        print(f"   Avg latency: {sum(latencies)/len(latencies):.2f}s")
        print(f"   P95 latency: {p95_latency:.2f}s")

    feedback_samples = [
        r.get("feedback", "")
        for r in success_results
        if isinstance(r.get("feedback", ""), str) and r.get("feedback", "").strip()
    ]
    if feedback_samples:
        print("   Sample feedback:")
        for idx, sample in enumerate(feedback_samples[:3], start=1):
            compact = " ".join(sample.split())
            print(f"     {idx}. {compact[:140]}...")


def evaluate_submission_with_prompt(submission_id: int, mentor: ADFMentor, submission_path: str, question: str, prompt: str):
    """Evaluate a single submission where prompt/question are required (e.g., .txt)."""
    global success_count, fail_count, total_score

    try:
        start_time = time.time()
        result = mentor.evaluate_adf(submission_path, question=question, prompt=prompt)
        elapsed = time.time() - start_time

        with results_lock:
            success_count += 1
            total_score += result['score']

        return {
            "id": submission_id,
            "status": "success",
            "score": result['score'],
            "time": elapsed,
            "feedback_length": len(result.get('feedback', '')),
            "feedback": result.get('feedback', '')
        }
    except Exception as e:
        with results_lock:
            fail_count += 1

        return {
            "id": submission_id,
            "status": "failed",
            "error": str(e),
            "time": 0
        }


def evaluate_submission(submission_id: int, mentor: ADFMentor, submission_path: str):
    """Evaluate a single submission with error handling."""
    global success_count, fail_count, total_score
    
    try:
        start_time = time.time()
        result = mentor.evaluate_adf(submission_path)
        elapsed = time.time() - start_time
        
        with results_lock:
            success_count += 1
            total_score += result['score']
        
        return {
            "id": submission_id,
            "status": "success",
            "score": result['score'],
            "time": elapsed,
            "feedback_length": len(result.get('feedback', '')),
            "feedback": result.get('feedback', '')
        }
    except Exception as e:
        with results_lock:
            fail_count += 1
        
        return {
            "id": submission_id,
            "status": "failed",
            "error": str(e),
            "time": 0
        }


def test_sequential_processing():
    """Test 1: Sequential processing (safe, slow)."""
    print("\n" + "="*70)
    print("TEST 1: SEQUENTIAL PROCESSING (100 submissions)")
    print("="*70)
    print("Processing one by one (SAFE but SLOW)...")
    
    global success_count, fail_count, total_score
    success_count = fail_count = total_score = 0
    
    mentor = create_mentor_or_report()
    if mentor is None:
        return []
    
    start_time = time.time()
    results = []
    
    for i in range(100):
        print(f"Processing submission {i+1}/100...", end="\r")
        result = evaluate_submission(i+1, mentor, "lesson-3.zip")
        results.append(result)
    
    elapsed = time.time() - start_time
    
    print(f"\n\n✅ Sequential Processing Complete!")
    print(f"   Time taken: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)")
    print(f"   Success: {success_count}/100")
    print(f"   Failed: {fail_count}/100")
    if success_count > 0:
        print(f"   Average score: {total_score/success_count:.1f}/100")
    print(f"   Average time per submission: {elapsed/100:.2f}s")
    
    return results


def test_concurrent_no_limits():
    """Test 2: Concurrent processing WITHOUT rate limiting (will fail)."""
    print("\n" + "="*70)
    print("TEST 2: CONCURRENT WITHOUT RATE LIMITING (100 at once)")
    print("="*70)
    print("⚠️  This will likely HIT RATE LIMITS and fail...")
    
    global success_count, fail_count, total_score
    success_count = fail_count = total_score = 0
    
    mentor = create_mentor_or_report()
    if mentor is None:
        return []
    
    start_time = time.time()
    results = []
    
    with ThreadPoolExecutor(max_workers=100) as executor:
        futures = [
            executor.submit(evaluate_submission, i+1, mentor, "lesson-3.zip")
            for i in range(100)
        ]
        
        for future in as_completed(futures):
            result = future.result()
            results.append(result)
            if result['status'] == 'success':
                print(f"✅ Submission {result['id']}: {result['score']}/100", end="\r")
            else:
                print(f"❌ Submission {result['id']}: FAILED - {result['error'][:50]}", end="\r")
    
    elapsed = time.time() - start_time
    
    print(f"\n\n⚠️  Concurrent (No Limits) Complete!")
    print(f"   Time taken: {elapsed:.2f} seconds")
    print(f"   Success: {success_count}/100")
    print(f"   Failed: {fail_count}/100 (likely rate limit errors)")
    if success_count > 0:
        print(f"   Average score: {total_score/success_count:.1f}/100")
    
    return results


def test_concurrent_with_rate_limiting():
    """Test 3: Concurrent with rate limiting (RECOMMENDED for PAID TIER)."""
    print("\n" + "="*70)
    print("TEST 3: CONCURRENT WITH RATE LIMITING (batches of 50)")
    print("="*70)
    print("Processing in controlled batches (OPTIMIZED FOR PAID TIER)...")
    print("Paid tier limits: 1,000 RPM, 4M TPM")
    
    global success_count, fail_count, total_score
    success_count = fail_count = total_score = 0
    
    mentor = create_mentor_or_report()
    if mentor is None:
        return []
    
    start_time = time.time()
    results = []
    
    # Process in batches of 50 with minimal delays (paid tier)
    batch_size = 50
    delay_between_batches = 5  # seconds (minimal delay for paid tier)
    
    num_batches = (100 + batch_size - 1) // batch_size  # 100 / 50 = 2 batches
    for batch_num in range(num_batches):
        batch_start = batch_num * batch_size
        batch_end = min((batch_num + 1) * batch_size, 100)
        batch_count = batch_end - batch_start
        
        print(f"\nBatch {batch_num+1}/{num_batches} starting ({batch_count} submissions)...")
        
        with ThreadPoolExecutor(max_workers=batch_count) as executor:
            futures = [
                executor.submit(
                    evaluate_submission, 
                    batch_start + i + 1, 
                    mentor, 
                    "lesson-3.zip"
                )
                for i in range(batch_count)
            ]
            
            for future in as_completed(futures):
                result = future.result()
                results.append(result)
                if result['status'] == 'success':
                    print(f"  ✅ Submission {result['id']}: {result['score']}/100")
                else:
                    print(f"  ❌ Submission {result['id']}: FAILED")
        
        if batch_num < num_batches - 1:  # Don't delay after last batch
            print(f"  Waiting {delay_between_batches}s before next batch...")
            time.sleep(delay_between_batches)
    
    elapsed = time.time() - start_time
    
    summarize_results(results, "Rate-Limited Concurrent Processing", elapsed)
    
    return results


def test_small_batch():
    """Quick test with 10 submissions to verify setup."""
    print("\n" + "="*70)
    print("QUICK TEST: 10 submissions (verify before full test)")
    print("="*70)
    
    global success_count, fail_count, total_score
    success_count = fail_count = total_score = 0
    
    mentor = create_mentor_or_report()
    if mentor is None:
        return False
    
    start_time = time.time()
    
    for i in range(10):
        print(f"Processing {i+1}/10...", end="\r")
        evaluate_submission(i+1, mentor, "lesson-3.zip")
    
    elapsed = time.time() - start_time
    
    print(f"\n\n✅ Quick Test Complete!")
    print(f"   Time: {elapsed:.2f}s")
    print(f"   Success: {success_count}/10")
    print(f"   Failed: {fail_count}/10")
    if success_count > 0:
        print(f"   Average score: {total_score/success_count:.1f}/100")
    
    if fail_count > 0:
        print(f"\n⚠️  WARNING: {fail_count} failures detected!")
        print("   Check API key and rate limits before running full test.")
        return False
    
    return True


def _create_invalid_zip(path: str):
    """Create an intentionally invalid ZIP (non-JSON content) for crash-resilience testing."""
    with tempfile.TemporaryDirectory() as tmp:
        bad_file = os.path.join(tmp, "bad_payload.xlsx")
        with open(bad_file, "wb") as f:
            f.write(b"not a valid adf json")
        with zipfile.ZipFile(path, "w") as z:
            z.write(bad_file, "bad_payload.xlsx")


def test_hard_mixed_inputs():
    """Test 6: Hard mixed-quality parallel run (valid + text + invalid)."""
    print("\n" + "="*70)
    print("TEST 6: HARD MIXED INPUTS (Crash-Resilience + Feedback)")
    print("="*70)
    print("Running 100 parallel evaluations with mixed submission types...")

    global success_count, fail_count, total_score
    success_count = fail_count = total_score = 0

    mentor = create_mentor_or_report()
    if mentor is None:
        return []

    invalid_zip_path = "test_invalid_mixed.zip"
    _create_invalid_zip(invalid_zip_path)

    text_question = "Explain Linked Services and handling timeout errors in ADF."
    text_prompt = (
        "Evaluate written ADF understanding for technical accuracy, clarity, and practical relevance. "
        "Return score 0-100 and concise feedback."
    )

    tasks = []
    submission_id = 1

    # 50 valid ZIP submissions
    for _ in range(50):
        tasks.append(("normal", submission_id, "lesson-3.zip"))
        submission_id += 1

    # 20 folder submissions
    for _ in range(20):
        tasks.append(("normal", submission_id, "lesson-3/lesson-3/"))
        submission_id += 1

    # 20 text-only submissions with prompt
    for _ in range(20):
        tasks.append(("text", submission_id, "lesson-3/lesson-3/answers.txt"))
        submission_id += 1

    # 10 intentionally invalid submissions (should return 0 but not crash)
    for _ in range(10):
        tasks.append(("normal", submission_id, invalid_zip_path))
        submission_id += 1

    start_time = time.time()
    results = []

    with ThreadPoolExecutor(max_workers=50) as executor:
        futures = []
        for task_type, sid, path in tasks:
            if task_type == "text":
                futures.append(
                    executor.submit(
                        evaluate_submission_with_prompt,
                        sid,
                        mentor,
                        path,
                        text_question,
                        text_prompt,
                    )
                )
            else:
                futures.append(executor.submit(evaluate_submission, sid, mentor, path))

        for future in as_completed(futures):
            result = future.result()
            results.append(result)
            if result["status"] == "success":
                print(f"  ✅ Submission {result['id']}: {result['score']}/100")
            else:
                print(f"  ❌ Submission {result['id']}: FAILED")

    elapsed = time.time() - start_time
    summarize_results(results, "Hard Mixed Inputs", elapsed)

    # Crash-resilience assertions as runtime report
    zero_score_count = sum(1 for r in results if r.get("status") == "success" and r.get("score") == 0)
    print(f"   Zero-score responses (expected includes invalid files): {zero_score_count}")
    print("   Pass condition: script completes with no unhandled exceptions.")

    if os.path.exists(invalid_zip_path):
        os.remove(invalid_zip_path)

    return results


def test_feedback_quality_audit():
    """Test 7: Validate feedback usefulness characteristics over repeated runs."""
    print("\n" + "="*70)
    print("TEST 7: FEEDBACK QUALITY AUDIT (30 submissions)")
    print("="*70)
    print("Checking feedback length and presence of actionable keywords...")

    global success_count, fail_count, total_score
    success_count = fail_count = total_score = 0

    mentor = create_mentor_or_report()
    if mentor is None:
        return []
    start_time = time.time()

    results = []
    with ThreadPoolExecutor(max_workers=15) as executor:
        futures = [
            executor.submit(evaluate_submission, i + 1, mentor, "lesson-3.zip")
            for i in range(30)
        ]
        for future in as_completed(futures):
            results.append(future.result())

    elapsed = time.time() - start_time
    summarize_results(results, "Feedback Quality Audit", elapsed)

    success_results = [r for r in results if r.get("status") == "success"]
    feedbacks = [r.get("feedback", "") for r in success_results if r.get("feedback")]
    if not feedbacks:
        print("   No feedback text captured in this mode.")
        return results

    keywords = ["parameter", "error", "pipeline", "activity", "security", "missing", "improve"]
    actionable_count = 0
    for fb in feedbacks:
        lower_fb = fb.lower()
        if any(k in lower_fb for k in keywords):
            actionable_count += 1

    avg_len = sum(len(fb) for fb in feedbacks) / len(feedbacks)
    print(f"   Feedback samples analyzed: {len(feedbacks)}")
    print(f"   Average feedback length: {avg_len:.1f} chars")
    print(f"   Actionable-keyword hit rate: {actionable_count}/{len(feedbacks)}")

    return results


def test_single_professional_report():
    """Test 8: Single submission with full professional report output."""
    print("\n" + "="*70)
    print("TEST 8: SINGLE PROFESSIONAL REPORT")
    print("="*70)

    mentor = create_mentor_or_report()
    if mentor is None:
        return None

    submission_path = "lesson-3.zip"
    question = "Evaluate this Azure Data Factory submission for quality, completeness, and best practices."

    start_time = time.time()
    result = mentor.evaluate_adf(submission_path, question=question)
    elapsed = time.time() - start_time

    print(f"Submission: {submission_path}")
    print(f"Elapsed: {elapsed:.2f}s")
    print(f"Score: {result.get('score')}/100")

    print("\n--- FEEDBACK ---")
    print(result.get("feedback", "<no feedback>"))

    print("\n--- METADATA REPORT ---")
    print(result.get("metadata_report", "<no metadata report>"))

    print("\n--- RAW RESULT JSON ---")
    print(json.dumps(result, indent=2, ensure_ascii=False))

    return result


if __name__ == "__main__":
    print("\n" + "="*70)
    print("ADFMentor LOAD TESTING - Production Simulation")
    print("="*70)
    print("\nThis simulates 100 students submitting homework simultaneously.")
    print("Using lesson-3.zip as the sample submission.")
    print("\n🎯 PAID TIER DETECTED: Optimized for 1,000 RPM limits\n")
    
    # Check if lesson-3.zip exists
    if not os.path.exists("lesson-3.zip"):
        print("❌ ERROR: lesson-3.zip not found!")
        print("   Please create it first:")
        print("   Compress-Archive -Path lesson-3 -DestinationPath lesson-3.zip")
        exit(1)
    
    print("Choose test mode:")
    print("1. Quick test (10 submissions) - Verify setup (~30s)")
    print("2. Sequential (100 submissions) - Safe but slow (~5 min)")
    print("3. Concurrent no limits (100 at once) - Fast but risky (~20s)")
    print("4. Concurrent with rate limiting (2 batches of 50) - RECOMMENDED (~15s)")
    print("5. All tests (will take ~6-7 minutes total)")
    print("6. Hard mixed inputs (100 mixed valid/invalid/text)")
    print("7. Feedback quality audit (30 runs)")
    print("8. Single professional report (1 submission, full output)")
    
    choice = input("\nEnter choice (1-8): ").strip()
    
    if choice == "1":
        test_small_batch()
    elif choice == "2":
        test_sequential_processing()
    elif choice == "3":
        test_concurrent_no_limits()
    elif choice == "4":
        test_concurrent_with_rate_limiting()
    elif choice == "5":
        if test_small_batch():
            print("\n✅ Quick test passed! Proceeding to full tests...")
            time.sleep(2)
            test_sequential_processing()
            time.sleep(5)
            test_concurrent_no_limits()
            time.sleep(5)
            test_concurrent_with_rate_limiting()
    elif choice == "6":
        test_hard_mixed_inputs()
    elif choice == "7":
        test_feedback_quality_audit()
    elif choice == "8":
        test_single_professional_report()
    else:
        print("Invalid choice!")
    
    print("\n" + "="*70)
    print("Load testing complete!")
    print("="*70)
