# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "datasets",
#     "huggingface_hub",
# ]
# ///
#!/usr/bin/env python3
"""
Agent Solve Rate Experiment for SWE-Bench Pro
Uses DeepSeek-V4-Flash (free HF inference) to attempt solving SWE-Bench Pro tasks.
Measures format-compliant patch generation rate (NOT actual correctness).
"""
import json
import time
import re
import sys
from pathlib import Path

# Configuration
NUM_TASKS = 20  # Test on 20 tasks for statistical significance
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
MAX_TOKENS = 2048
RATE_LIMIT_DELAY = 2.0  # seconds between calls to avoid rate limits
OUTPUT_FILE = "/tmp/agent_solve_results.json"

def load_dataset():
    """Load SWE-Bench Pro dataset from HuggingFace."""
    from datasets import load_dataset
    ds = load_dataset("ScaleAI/SWE-bench_Pro", split="test")
    return ds

def create_prompt(instance):
    """Create a prompt for the model to generate a patch."""
    repo = instance.get("repo", "unknown")
    instance_id = instance.get("instance_id", "unknown")
    problem_statement = instance.get("problem_statement", "")
    base_commit = instance.get("base_commit", "")
    
    prompt = f"""You are an expert software engineer. Given the following issue in the repository {repo}, generate a patch to fix the issue.

Issue: {problem_statement}

Please provide a unified diff patch that fixes this issue. The patch should:
1. Be in unified diff format (--- a/file.py, +++ b/file.py)
2. Only modify the necessary files
3. Be minimal and focused on the fix

Output ONLY the patch in unified diff format, no explanation:"""
    
    return prompt

def call_model(prompt, max_retries=3):
    """Call DeepSeek-V4-Flash with retry logic."""
    from huggingface_hub import InferenceClient
    
    client = InferenceClient()
    
    for attempt in range(max_retries):
        try:
            response = client.chat.completions.create(
                model=MODEL,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=MAX_TOKENS,
                temperature=0.0
            )
            return response.choices[0].message.content
        except Exception as e:
            if "rate" in str(e).lower() and attempt < max_retries - 1:
                wait = RATE_LIMIT_DELAY * (attempt + 1)
                print(f"Rate limited, waiting {wait}s...")
                time.sleep(wait)
            else:
                print(f"Error calling model: {e}")
                return None
    return None

def is_valid_patch(response):
    """Check if the response looks like a valid unified diff patch."""
    if not response:
        return False, "No response"
    
    # Check for unified diff markers
    has_diff_header = bool(re.search(r'^diff --git', response, re.MULTILINE) or 
                          re.search(r'^---', response, re.MULTILINE) or
                          re.search(r'^\+\+\+', response, re.MULTILINE))
    has_hunk_header = bool(re.search(r'^@@', response, re.MULTILINE))
    has_additions = bool(re.search(r'^\+[^+]', response, re.MULTILINE))
    has_deletions = bool(re.search(r'^-[^-]', response, re.MULTILINE))
    
    if has_diff_header and has_hunk_header and (has_additions or has_deletions):
        return True, "Valid unified diff"
    elif has_hunk_header:
        return True, "Has hunk headers"
    elif has_additions or has_deletions:
        return True, "Has changes"
    else:
        return False, "No diff markers found"

def main():
    print(f"Loading SWE-Bench Pro dataset...")
    ds = load_dataset()
    print(f"Total instances: {len(ds)}")
    
    # Sample tasks
    import random
    random.seed(42)
    indices = random.sample(range(len(ds)), min(NUM_TASKS, len(ds)))
    tasks = [ds[i] for i in indices]
    
    results = []
    success_count = 0
    error_count = 0
    
    for i, instance in enumerate(tasks):
        instance_id = instance.get("instance_id", f"task_{i}")
        print(f"\n[{i+1}/{len(tasks)}] Processing {instance_id}...")
        
        prompt = create_prompt(instance)
        response = call_model(prompt)
        
        is_valid, reason = is_valid_patch(response)
        
        result = {
            "instance_id": instance_id,
            "repo": instance.get("repo", ""),
            "response_length": len(response) if response else 0,
            "is_valid_patch": is_valid,
            "validation_reason": reason,
            "response_preview": response[:500] if response else ""
        }
        results.append(result)
        
        if is_valid:
            success_count += 1
            print(f"  ✓ Valid patch ({reason})")
        else:
            error_count += 1
            print(f"  ✗ {reason}")
        
        # Rate limiting
        time.sleep(RATE_LIMIT_DELAY)
    
    # Summary
    summary = {
        "model": MODEL,
        "total_tasks": len(tasks),
        "valid_patches": success_count,
        "invalid_patches": error_count,
        "format_compliance_rate": success_count / len(tasks) if tasks else 0,
        "results": results
    }
    
    # Save results
    with open(OUTPUT_FILE, "w") as f:
        json.dump(summary, f, indent=2)
    
    print(f"\n{'='*60}")
    print(f"RESULTS SUMMARY")
    print(f"{'='*60}")
    print(f"Model: {MODEL}")
    print(f"Tasks tested: {len(tasks)}")
    print(f"Valid patches: {success_count}/{len(tasks)} ({success_count/len(tasks)*100:.1f}%)")
    print(f"Results saved to: {OUTPUT_FILE}")
    
    return summary

if __name__ == "__main__":
    main()