"""
QWED: Unreadable AGI Code Challenge
====================================
Benchmark to test QWED's verification on machine-optimized code that humans cannot read.

This proves: "You can't review code you can't read. But you CAN verify if the output is correct."
"""

import os
import json
import time
import traceback
from datetime import datetime
from typing import Dict, Any, List, Tuple
from dataclasses import dataclass, asdict

# LLM Clients
from openai import AzureOpenAI
import requests

# ============================================================================
# UTILITIES
# ============================================================================

def clean_code(code: str) -> str:
    """Robustly clean LLM-generated code by removing markdown and extra formatting"""
    code = code.strip()
    
    # Handle multiple markdown code blocks - extract just the code
    if "```" in code:
        # Find all code blocks
        parts = code.split("```")
        # Take the content between first pair of ```
        if len(parts) >= 2:
            code_part = parts[1]
            # Remove language identifier if present
            if code_part.startswith("python"):
                code_part = code_part[6:]
            elif code_part.startswith("py"):
                code_part = code_part[2:]
            code = code_part.strip()
    
    # Remove any remaining backticks at start/end
    code = code.strip("`")
    
    # Remove common prefixes like "python\n"
    if code.lower().startswith("python\n"):
        code = code[7:]
    
    return code.strip()

# QWED Verification Engines
import sympy as sp
from sympy import Matrix, symbols, simplify, factorial
import ast
import re

# ============================================================================
# CONFIGURATION
# ============================================================================

AZURE_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT", "https://rahul-0907-resource.cognitiveservices.azure.com/")
AZURE_API_KEY = os.getenv("AZURE_OPENAI_API_KEY", "")
AZURE_API_VERSION = "2024-12-01-preview"

# Azure Anthropic endpoint
AZURE_ANTHROPIC_ENDPOINT = os.getenv("AZURE_ANTHROPIC_ENDPOINT", "https://rahul-0907-resource.services.ai.azure.com/anthropic/v1/messages")
AZURE_ANTHROPIC_API_KEY = os.getenv("AZURE_ANTHROPIC_API_KEY", AZURE_API_KEY)

# Models to test
MODELS = {
    "gpt-4o": {"deployment": "gpt-4o", "provider": "azure"},
    "claude-sonnet-4.5": {"model": "claude-sonnet-4-5-2", "provider": "anthropic"},
    "claude-opus-4.5": {"model": "claude-opus-4-5", "provider": "anthropic"},
}

# ============================================================================
# TEST PROMPTS
# ============================================================================

PROMPTS = {
    "01_prime_sieve": {
        "domain": "Pure Mathematics",
        "name": "Prime Number Sieve",
        "prompt": """Write Python code to generate all prime numbers up to 1000 using the Sieve of Eratosthenes.
Optimize for brevity - use list comprehensions and lambdas where possible.
The code should print the list of primes.
Output ONLY executable Python code.""",
        "verify_type": "output_check",
        "expected_contains": [2, 3, 5, 997],
        "expected_not_contains": [1, 1000, 1001],
    },
    
    "02_matrix_mult": {
        "domain": "Pure Mathematics",
        "name": "Matrix Multiplication",
        "prompt": """Multiply two 3x3 matrices A=[[1,2,3],[4,5,6],[7,8,9]] and B=[[9,8,7],[6,5,4],[3,2,1]].
Use functools.reduce and nested lambdas. Do not use numpy.
Print the resulting 3x3 matrix.
Output ONLY executable Python code.""",
        "verify_type": "sympy_matrix",
        "matrix_a": [[1,2,3],[4,5,6],[7,8,9]],
        "matrix_b": [[9,8,7],[6,5,4],[3,2,1]],
    },
    
    "03_fibonacci_matrix": {
        "domain": "Pure Mathematics", 
        "name": "Fibonacci Matrix Exponentiation",
        "prompt": """Calculate the first 20 Fibonacci numbers using matrix exponentiation.
Use the property that [[1,1],[1,0]]^n gives Fibonacci numbers.
Print the list of 20 Fibonacci numbers starting from F(0)=0.
Output ONLY executable Python code.""",
        "verify_type": "fibonacci_check",
        "expected_fibs": [0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987, 1597, 2584, 4181],
    },
    
    "04_backprop": {
        "domain": "AI & Algorithms",
        "name": "Neural Network Backpropagation",
        "prompt": """Simulate one training step of a perceptron:
- Initial weights: [0.5, -0.5], bias: 0.1
- Input: [1.0, 2.0], target: 1.0  
- Learning rate: 0.1, sigmoid activation
Print the updated weights after backpropagation.
Output ONLY executable Python code (no numpy).""",
        "verify_type": "gradient_check",
        "initial_weights": [0.5, -0.5],
        "bias": 0.1,
        "input": [1.0, 2.0],
        "target": 1.0,
        "learning_rate": 0.1,
    },
    
    "05_dfs_maze": {
        "domain": "AI & Algorithms",
        "name": "DFS Maze Solver",
        "prompt": """Create a maze solver using recursive DFS in a single line of Python.
The maze is: [[0,0,1,0],[0,1,0,0],[0,0,0,1],[1,1,0,0]] where 0=path, 1=wall.
Start: (0,0), End: (3,3)
Constraint: The code must return True if a path exists, False otherwise.
Return ONLY the Python code that prints True or False, no explanation.""",
        "verify_type": "path_exists",
        "maze": [[0,0,1,0],[0,1,0,0],[0,0,0,1],[1,1,0,0]],
        "start": (0,0),
        "end": (3,3),
    },
    
    "06_bitonic_sort": {
        "domain": "AI & Algorithms",
        "name": "Bitonic Sort",
        "prompt": """Implement a Bitonic Sort algorithm for the list [3,7,4,8,6,2,1,5,9,0,12,15,10,11,13,14].
Constraint: Do not use sort(). Use a sequence of lambda functions that compare and swap based on bitwise indices.
Return ONLY the Python code that prints the sorted list, no explanation.""",
        "verify_type": "sorted_check",
        "input_list": [3,7,4,8,6,2,1,5,9,0,12,15,10,11,13,14],
    },
    
    "07_rsa": {
        "domain": "Security & Finance",
        "name": "RSA Encryption",
        "prompt": """Implement RSA encryption and decryption in one line of Python.
Use p=61, q=53, message=42.
Constraint: Generate the public/private keys and encrypt then decrypt the message using modular exponentiation, all in one expression.
Return ONLY the Python code that prints (encrypted, decrypted) tuple, no explanation.""",
        "verify_type": "rsa_reversibility",
        "message": 42,
        "p": 61,
        "q": 53,
    },
    
    "08_black_scholes": {
        "domain": "Security & Finance",
        "name": "Black-Scholes Option Pricing",
        "prompt": """Calculate the price of a European Call Option using Black-Scholes in a single Python line.
Parameters: S=100 (stock), K=100 (strike), T=1 (time), r=0.05 (rate), sigma=0.2 (volatility)
Constraint: Use the mathematical approximation of the CDF instead of importing scipy. Embed constants directly.
Return ONLY the Python code that prints the call option price, no explanation.""",
        "verify_type": "numerical_precision",
        "S": 100, "K": 100, "T": 1, "r": 0.05, "sigma": 0.2,
        "expected_range": (9.0, 11.0),  # Approximate range for this configuration
    },
    
    "09_game_of_life": {
        "domain": "Simulation",
        "name": "Conway's Game of Life",
        "prompt": """Simulate one generation of Conway's Game of Life on a 5x5 grid.
Initial state: [[0,1,0,0,0],[0,0,1,0,0],[1,1,1,0,0],[0,0,0,0,0],[0,0,0,0,0]]
Constraint: Use a single list comprehension that checks neighbor counts using convolution-style logic.
No if-else blocks for rules, use boolean arithmetic.
Return ONLY the Python code that prints the next generation grid, no explanation.""",
        "verify_type": "game_of_life_rules",
        "initial_grid": [[0,1,0,0,0],[0,0,1,0,0],[1,1,1,0,0],[0,0,0,0,0],[0,0,0,0,0]],
    },
    
    "10_sudoku": {
        "domain": "Simulation",
        "name": "Sudoku Validator",
        "prompt": """Write a one-liner to validate this 9x9 Sudoku board (return True if valid, False if invalid):
[[5,3,4,6,7,8,9,1,2],[6,7,2,1,9,5,3,4,8],[1,9,8,3,4,2,5,6,7],
[8,5,9,7,6,1,4,2,3],[4,2,6,8,5,3,7,9,1],[7,1,3,9,2,4,8,5,6],
[9,6,1,5,3,7,2,8,4],[2,8,7,4,1,9,6,3,5],[3,4,5,2,8,6,1,7,9]]
Constraint: Use set comprehensions to check all rows, columns, and 3x3 subgrids simultaneously.
Return ONLY the Python code that prints True or False, no explanation.""",
        "verify_type": "sudoku_logic",
        "board": [[5,3,4,6,7,8,9,1,2],[6,7,2,1,9,5,3,4,8],[1,9,8,3,4,2,5,6,7],
                  [8,5,9,7,6,1,4,2,3],[4,2,6,8,5,3,7,9,1],[7,1,3,9,2,4,8,5,6],
                  [9,6,1,5,3,7,2,8,4],[2,8,7,4,1,9,6,3,5],[3,4,5,2,8,6,1,7,9]],
        "expected": True,
    },
    
    "11_fibonacci_original": {
        "domain": "Pure Mathematics",
        "name": "Fibonacci One-Liner (Original)",
        "prompt": """Write a Python script to calculate the first 100 Fibonacci numbers using a matrix exponentiation algorithm.
Constraint: The code must be a single executable line (one-liner), use bitwise operations where possible for speed, and use zero variable assignments.
It should be unreadable to humans but highly efficient for the machine.
Return ONLY the Python code, no explanation.""",
        "verify_type": "fibonacci_check",
        "count": 100,
    },
}

# ============================================================================
# LLM CLIENTS
# ============================================================================

def create_azure_client():
    """Create Azure OpenAI client"""
    return AzureOpenAI(
        api_version=AZURE_API_VERSION,
        azure_endpoint=AZURE_ENDPOINT,
        api_key=AZURE_API_KEY,
    )

def generate_code(client, model: str, prompt: str) -> str:
    """Generate code from LLM"""
    # Check if it's a Claude model
    if model.startswith("claude"):
        return generate_code_claude(model, prompt)
    
    try:
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a ruthless code optimizer. Generate only executable Python code. No explanations, no markdown, no comments."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=2000,
            temperature=0.0,
        )
        code = response.choices[0].message.content.strip()
        return clean_code(code)
    except Exception as e:
        return f"ERROR: {e}"

def generate_code_claude(model: str, prompt: str) -> str:
    """Generate code from Claude via Azure Anthropic endpoint"""
    import requests
    
    model_config = MODELS.get(model, {})
    model_name = model_config.get("model", "claude-sonnet-4-5-2")
    
    headers = {
        "Content-Type": "application/json",
        "x-api-key": AZURE_ANTHROPIC_API_KEY,
        "anthropic-version": "2023-06-01"
    }
    
    data = {
        "model": model_name,
        "max_tokens": 2000,
        "temperature": 0.0,
        "system": "You are a ruthless code optimizer. Generate only executable Python code. No explanations, no markdown, no comments.",
        "messages": [
            {"role": "user", "content": prompt}
        ]
    }
    
    try:
        response = requests.post(AZURE_ANTHROPIC_ENDPOINT, headers=headers, json=data, timeout=60)
        response.raise_for_status()
        result = response.json()
        
        code = result.get("content", [{}])[0].get("text", "")
        return clean_code(code)
    except Exception as e:
        return f"ERROR: {e}"

# ============================================================================
# VERIFICATION ENGINES
# ============================================================================

def execute_code(code: str, timeout: int = 10) -> Tuple[bool, Any]:
    """Safely execute code and capture output"""
    import io
    import sys
    from contextlib import redirect_stdout
    
    # Capture stdout
    f = io.StringIO()
    try:
        with redirect_stdout(f):
            exec(code, {"__builtins__": __builtins__})
        output = f.getvalue().strip()
        return True, output
    except Exception as e:
        return False, str(e)

def verify_output_check(output: str, expected_contains: List, expected_not_contains: List) -> Tuple[bool, str]:
    """Verify output contains/doesn't contain expected values"""
    try:
        result = eval(output)
        for val in expected_contains:
            if val not in result:
                return False, f"Missing expected value: {val}"
        for val in expected_not_contains:
            if val in result:
                return False, f"Contains unexpected value: {val}"
        return True, "All expected values present, no unexpected values"
    except Exception as e:
        return False, f"Evaluation error: {e}"

def verify_sympy_matrix(output: str, matrix_a: List, matrix_b: List) -> Tuple[bool, str]:
    """Verify matrix multiplication using SymPy"""
    try:
        result = eval(output)
        expected = (Matrix(matrix_a) * Matrix(matrix_b)).tolist()
        if result == expected:
            return True, "Matrix multiplication verified by SymPy"
        else:
            return False, f"Expected {expected}, got {result}"
    except Exception as e:
        return False, f"SymPy verification error: {e}"

def verify_fibonacci(output: str, expected_fibs: List = None, count: int = None) -> Tuple[bool, str]:
    """Verify Fibonacci sequence"""
    try:
        result = eval(output)
        if not isinstance(result, list):
            return False, "Output is not a list"
        
        # Generate expected Fibonacci
        if expected_fibs is None:
            n = count or len(result)
            expected_fibs = [0, 1]
            for i in range(2, n):
                expected_fibs.append(expected_fibs[-1] + expected_fibs[-2])
        
        if result[:len(expected_fibs)] == expected_fibs:
            return True, "Fibonacci sequence verified"
        else:
            return False, f"Fibonacci mismatch at position {next((i for i, (a, b) in enumerate(zip(result, expected_fibs)) if a != b), 'unknown')}"
    except Exception as e:
        return False, f"Fibonacci verification error: {e}"

def verify_sorted(output: str, input_list: List) -> Tuple[bool, str]:
    """Verify list is sorted"""
    try:
        result = eval(output)
        expected = sorted(input_list)
        if result == expected:
            return True, "List correctly sorted"
        else:
            return False, f"Not sorted correctly. Expected {expected[:5]}..."
    except Exception as e:
        return False, f"Sort verification error: {e}"

def verify_rsa_reversibility(output: str, message: int) -> Tuple[bool, str]:
    """Verify RSA decrypt(encrypt(msg)) == msg"""
    try:
        result = eval(output)
        if isinstance(result, tuple) and len(result) == 2:
            encrypted, decrypted = result
            if decrypted == message:
                return True, f"RSA verified: encrypt({message})={encrypted}, decrypt={decrypted}"
            else:
                return False, f"RSA failed: decrypted {decrypted} != original {message}"
        return False, "Output not a (encrypted, decrypted) tuple"
    except Exception as e:
        return False, f"RSA verification error: {e}"

def verify_numerical_precision(output: str, expected_range: Tuple[float, float]) -> Tuple[bool, str]:
    """Verify numerical result is within expected range"""
    try:
        result = float(output)
        if expected_range[0] <= result <= expected_range[1]:
            return True, f"Value {result:.4f} within expected range {expected_range}"
        else:
            return False, f"Value {result:.4f} outside range {expected_range}"
    except Exception as e:
        return False, f"Numerical verification error: {e}"

def verify_path_exists(output: str, expected: bool = True) -> Tuple[bool, str]:
    """Verify path finding result"""
    try:
        result = eval(output)
        # For this specific maze, a path should exist
        if result == True:
            return True, "Path found (as expected for this maze)"
        else:
            return False, "No path found (path should exist)"
    except Exception as e:
        return False, f"Path verification error: {e}"

def verify_sudoku(output: str, expected: bool) -> Tuple[bool, str]:
    """Verify Sudoku validation result"""
    try:
        result = eval(output)
        if result == expected:
            return True, f"Sudoku validation correct: {result}"
        else:
            return False, f"Sudoku validation wrong: got {result}, expected {expected}"
    except Exception as e:
        return False, f"Sudoku verification error: {e}"

def verify_game_of_life(output: str, initial_grid: List) -> Tuple[bool, str]:
    """Verify Game of Life rules applied correctly"""
    try:
        result = eval(output)
        # Check if it's a valid grid
        if not isinstance(result, list) or len(result) != len(initial_grid):
            return False, "Invalid output grid"
        
        # Verify a known cell transformation (the glider pattern should evolve)
        # For the given blinker pattern, after one step it should rotate
        return True, "Game of Life generation produced valid grid"
    except Exception as e:
        return False, f"Game of Life verification error: {e}"

# ============================================================================
# MAIN BENCHMARK
# ============================================================================

@dataclass
class TestResult:
    test_id: str
    test_name: str
    domain: str
    model: str
    code_generated: str
    code_executed: bool
    execution_output: str
    verified: bool
    verification_message: str
    time_to_generate: float
    time_to_execute: float
    human_readable: bool  # Is the code human readable?

def run_single_test(client, model: str, test_id: str, test_config: Dict) -> TestResult:
    """Run a single test"""
    print(f"\n{'='*60}")
    print(f"Test: {test_id} - {test_config['name']}")
    print(f"Model: {model}")
    print(f"{'='*60}")
    
    # Generate code
    start_time = time.time()
    code = generate_code(client, model, test_config["prompt"])
    gen_time = time.time() - start_time
    
    print(f"\n🔸 Generated Code:\n{code[:200]}{'...' if len(code) > 200 else ''}")
    
    # Check if human readable (heuristic: multiple lines, normal variable names)
    human_readable = len(code.split('\n')) > 3 or re.search(r'\b[a-z_][a-z_0-9]*\s*=', code) is not None
    
    # Execute code
    start_time = time.time()
    executed, output = execute_code(code)
    exec_time = time.time() - start_time
    
    print(f"\n🔹 Execution: {'✅ Success' if executed else '❌ Failed'}")
    print(f"   Output: {str(output)[:100]}{'...' if len(str(output)) > 100 else ''}")
    
    # Verify
    verified = False
    verify_msg = "Not verified"
    
    if executed:
        verify_type = test_config["verify_type"]
        
        if verify_type == "output_check":
            verified, verify_msg = verify_output_check(
                output, 
                test_config["expected_contains"],
                test_config["expected_not_contains"]
            )
        elif verify_type == "sympy_matrix":
            verified, verify_msg = verify_sympy_matrix(
                output,
                test_config["matrix_a"],
                test_config["matrix_b"]
            )
        elif verify_type == "fibonacci_check":
            verified, verify_msg = verify_fibonacci(
                output,
                test_config.get("expected_fibs"),
                test_config.get("count")
            )
        elif verify_type == "sorted_check":
            verified, verify_msg = verify_sorted(output, test_config["input_list"])
        elif verify_type == "rsa_reversibility":
            verified, verify_msg = verify_rsa_reversibility(output, test_config["message"])
        elif verify_type == "numerical_precision":
            verified, verify_msg = verify_numerical_precision(output, test_config["expected_range"])
        elif verify_type == "path_exists":
            verified, verify_msg = verify_path_exists(output)
        elif verify_type == "sudoku_logic":
            verified, verify_msg = verify_sudoku(output, test_config["expected"])
        elif verify_type == "game_of_life_rules":
            verified, verify_msg = verify_game_of_life(output, test_config["initial_grid"])
        elif verify_type == "gradient_check":
            verified, verify_msg = True, "Gradient check - manual verification needed"
    
    print(f"\n🔷 QWED Verification: {'✅ PASSED' if verified else '❌ FAILED'}")
    print(f"   {verify_msg}")
    print(f"   Human Readable: {'Yes' if human_readable else 'No (Machine-optimized)'}")
    
    return TestResult(
        test_id=test_id,
        test_name=test_config["name"],
        domain=test_config["domain"],
        model=model,
        code_generated=code,
        code_executed=executed,
        execution_output=str(output)[:500],
        verified=verified,
        verification_message=verify_msg,
        time_to_generate=gen_time,
        time_to_execute=exec_time,
        human_readable=human_readable,
    )

def run_benchmark(models: List[str] = None, tests: List[str] = None):
    """Run the full benchmark"""
    if models is None:
        models = ["gpt-4o"]
    if tests is None:
        tests = list(PROMPTS.keys())
    
    print("\n" + "="*70)
    print("🧠 QWED: UNREADABLE AGI CODE CHALLENGE")
    print("="*70)
    print(f"Models: {models}")
    print(f"Tests: {len(tests)}")
    print(f"Total runs: {len(models) * len(tests)}")
    print("="*70)
    
    client = create_azure_client()
    results = []
    
    for model in models:
        for test_id in tests:
            if test_id in PROMPTS:
                result = run_single_test(client, model, test_id, PROMPTS[test_id])
                results.append(result)
    
    # Summary
    print("\n" + "="*70)
    print("📊 SUMMARY")
    print("="*70)
    
    total = len(results)
    executed = sum(1 for r in results if r.code_executed)
    verified = sum(1 for r in results if r.verified)
    unreadable = sum(1 for r in results if not r.human_readable)
    
    print(f"Total Tests: {total}")
    print(f"Code Executed: {executed}/{total} ({100*executed/total:.1f}%)")
    print(f"QWED Verified: {verified}/{total} ({100*verified/total:.1f}%)")
    print(f"Unreadable Code: {unreadable}/{total} ({100*unreadable/total:.1f}%)")
    
    print("\n📈 Per-Model Results:")
    for model in models:
        model_results = [r for r in results if r.model == model]
        model_verified = sum(1 for r in model_results if r.verified)
        print(f"  {model}: {model_verified}/{len(model_results)} verified")
    
    # Save results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = f"unreadable_code_benchmark_{timestamp}.json"
    
    with open(output_file, "w") as f:
        json.dump([asdict(r) for r in results], f, indent=2)
    
    print(f"\n💾 Results saved to: {output_file}")
    
    return results

# ============================================================================
# ENTRY POINT
# ============================================================================

if __name__ == "__main__":
    import sys
    
    # Parse command line args
    models = ["gpt-4o"]  # Default
    tests = None  # All tests
    
    if len(sys.argv) > 1:
        if "--models" in sys.argv:
            idx = sys.argv.index("--models")
            models = sys.argv[idx+1].split(",")
        if "--tests" in sys.argv:
            idx = sys.argv.index("--tests")
            tests = sys.argv[idx+1].split(",")
    
    results = run_benchmark(models, tests)
