import os
import json
import time
import requests
from typing import Dict, List
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm


def get_sandbox_result(code: str, cluster_key: str) -> Dict:
    """Submit Lean code to sandbox and get verification result."""
    data = {
        "clusterKey": cluster_key,
        "requestList": [{
            "code": code,
            "language": "lean",
            "run_timeout": 300,
        }]
    }
    headers = {"Content-Type": "application/json"}

    # Submit code for verification
    response = requests.post(
        "https://xxxx/run_batch_code_async",
        data=json.dumps(data),
        headers=headers,
        timeout=300
    )
    task_id = response.json()["data"]

    # Wait for completion
    while True:
        response = requests.get(
            f"https://xxxx/run_batch_code_await?taskId={task_id}&clusterKey={cluster_key}"
        )
        result = response.json()["data"]
        if result["requests"][0]["status"] == "success":
            break
        time.sleep(1)

    # Parse and return result
    run_result = json.loads(result["requests"][0]["result"]["run_result"]["stdout"])
    # print(run_result)
    return {
        "complete": run_result["complete"],
        "error": result["requests"][0]["result"]["run_result"].get("stderr", ""),
        "info": run_result,
        "verified_code": result["requests"][0]['request']['code']
    }


def batch_verify_lean_proofs(lean_contents: List[str], cluster_key: str, concurrency: int = 30) -> List[Dict]:
    """Verify a batch of Lean proofs using the API."""
    results = []

    with ThreadPoolExecutor(max_workers=concurrency) as executor:
        # Submit all verification tasks
        future_to_index = {
            executor.submit(get_sandbox_result, code, cluster_key): i
            for i, code in enumerate(lean_contents)
        }

        # Process results as they complete
        results = [None] * len(lean_contents)
        for future in tqdm(as_completed(future_to_index), total=len(lean_contents), desc="Verifying proofs"):
            index = future_to_index[future]
            try:
                results[index] = future.result()
            except Exception as e:
                results[index] = {"complete": False, "error": str(e)}

    return results
