"""
Single-instance smoke test for cascade validation.

This tests the ENTIRE pipeline end-to-end:
  1. Clone repo
  2. Set up conda environment
  3. Run cascade agent (T1 Llama-3.1-8B → T2 Llama-3.3-70B via HF Inference)
  4. Apply patch + test_patch
  5. Run FAIL_TO_PASS tests

This is the minimal test to prove the cascade works causally,
not just correlatively from trace simulation.

Submits to trackio for monitoring.

Usage (hf_jobs):
  operation: run
  script: "smoke_test_cascade.py"
  dependencies: ["huggingface_hub", "datasets", "trackio"]
  hardware: a10g-largex2 (need GPU for inference + CPU for conda + memory for clone)
  timeout: 4h
"""

import json
import os
import re
import subprocess
import sys
import tempfile
import time
import traceback
from datetime import datetime
from pathlib import Path

import datasets
import huggingface_hub
import trackio

# ============================================================
# Trackio setup
# ============================================================
trackio.init(
    project=os.environ.get("TRACKIO_PROJECT", "aco-smoke-test"),
)

# ============================================================
# CONFIG
# ============================================================
# The easiest SWE-bench instance: django bug with clear fix
INSTANCE_ID = os.environ.get("INSTANCE_ID", "django__django-14315")

T1_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
T2_MODEL = "meta-llama/Llama-3.3-70B-Instruct"

trackio.log_params({
    "instance_id": INSTANCE_ID,
    "t1_model": T1_MODEL,
    "t2_model": T2_MODEL,
    "strategy": "cascade",
})


def run_shell(cmd, cwd=None, timeout=120):
    result = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, timeout=timeout, shell=True)
    return result.returncode, result.stdout, result.stderr


def call_model(client, messages, max_tokens=4096):
    try:
        completion = client.chat.completions.create(
            model=client.model,
            messages=messages,
            max_tokens=max_tokens,
            temperature=0.2,
        )
        text = completion.choices[0].message.content
        itok = completion.usage.prompt_tokens if hasattr(completion, 'usage') and completion.usage else 0
        otok = completion.usage.completion_tokens if hasattr(completion, 'usage') and completion.usage else len(text) // 4
        return text, itok, otok
    except Exception as e:
        return f"[ERROR: {e}]", 0, 0


def extract_patch(text):
    for tag in ['patch', 'diff']:
        m = re.search(rf'<{tag}>(.*?)</{tag}>', text, re.DOTALL)
        if m:
            return m.group(1).strip()
    for block in ['diff', 'patch']:
        m = re.search(rf'```{block}\s*\n(.*?)```', text, re.DOTALL)
        if m:
            return m.group(1).strip()
    diff_match = re.search(r'(diff --git a/.*?(?:\n(?:@@|\+\+\+|diff --git|```|</).*)*)', text, re.DOTALL)
    if diff_match:
        return diff_match.group(1).strip()
    return None


def run_cascade_agent(instance, repo_dir):
    """Run T1 then T2. Maximum 30 turns each."""
    
    problem = instance.get("problem_statement", "")
    
    system = f"""You are fixing a bug in {instance['repo']}. Repository is at {repo_dir}.

Output format:
- Bash commands: <bash>command</bash>
- Final patch: <patch>your diff here</patch>
- Done: <submit>Done</submit>

First explore the codebase to understand the issue, then make a minimal fix and verify it."""

    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": f"PROBLEM:\n{problem}\n\nStart by exploring the repository."}
    ]
    
    tiers = [
        ("T1", T1_MODEL, 30),
        ("T2", T2_MODEL, 30),
    ]
    
    for tier_name, model_id, max_turns in tiers:
        print(f"\n{'='*50}")
        print(f"[{tier_name}] {model_id}")
        print(f"{'='*50}")
        
        trackio.log({"event": "tier_start", "tier": tier_name, "model": model_id})
        
        client = huggingface_hub.InferenceClient(model_id)
        total_itok = 0
        total_otok = 0
        
        for turn in range(max_turns):
            text, itok, otok = call_model(client, messages, max_tokens=4096)
            total_itok += itok
            total_otok += otok
            messages.append({"role": "assistant", "content": text})
            
            print(f"  Turn {turn+1}: {itok}+{otok} tokens, {len(text)} chars")
            
            patch = extract_patch(text)
            if patch:
                print(f"  ✅ PATCH FOUND ({len(patch)} chars)")
                trackio.log({
                    "event": "patch_found",
                    "tier": tier_name,
                    "turn": turn + 1,
                    "input_tokens": total_itok,
                    "output_tokens": total_otok,
                })
                return {
                    "patch": patch,
                    "tier": tier_name,
                    "turns": turn + 1,
                    "input_tokens": total_itok,
                    "output_tokens": total_otok,
                }
            
            cmds = re.findall(r'<bash>(.*?)</bash>', text, re.DOTALL)
            for cmd in cmds:
                cmd = cmd.strip()
                print(f"  $ {cmd[:120]}")
                rc, stdout, stderr = run_shell(cmd, cwd=str(repo_dir), timeout=30)
                output = (stdout + stderr)[:1500]
                if rc != 0:
                    output += f"\n[EXIT:{rc}]"
                messages.append({"role": "user", "content": f"<output>\n{output}\n</output>"})
            
            if "<submit>" in text:
                print(f"  [Submit without patch]")
                break
        
        trackio.log({"event": "tier_exhausted", "tier": tier_name, "turns": max_turns})
    
    return {"patch": None, "tier": None, "turns": 0, "input_tokens": 0, "output_tokens": 0}


def verify_patch(instance, model_patch, repo_dir, env_name=None):
    """Apply patches, run FAIL_TO_PASS tests."""
    base_commit = instance.get("base_commit", "")
    test_patch = instance.get("test_patch", "")
    f2p = instance.get("FAIL_TO_PASS", [])
    
    # Reset repo
    run_shell(f"cd {repo_dir} && git checkout -f {base_commit}", timeout=30)
    run_shell(f"cd {repo_dir} && git clean -fd", timeout=30)
    
    # Apply model patch
    patch_file = Path(repo_dir) / "_aco.patch"
    patch_file.write_text(model_patch)
    rc, out, err = run_shell(f"cd {repo_dir} && git apply --check _aco.patch", timeout=10)
    if rc != 0:
        return {"resolved": False, "error": f"model patch --check: {err[:300]}"}
    rc, out, err = run_shell(f"cd {repo_dir} && git apply _aco.patch", timeout=10)
    if rc != 0:
        return {"resolved": False, "error": f"model patch apply: {err[:300]}"}
    
    # Apply test_patch
    test_file = Path(repo_dir) / "_aco_test.patch"
    test_file.write_text(test_patch)
    rc, out, err = run_shell(f"cd {repo_dir} && git apply --check _aco_test.patch", timeout=10)
    if rc == 0:
        run_shell(f"cd {repo_dir} && git apply _aco_test.patch", timeout=10)
    else:
        run_shell(f"cd {repo_dir} && git apply --reject _aco_test.patch 2>/dev/null; true", timeout=10)
    
    # Run F2P tests
    cmd_prefix = f"conda run -n {env_name} " if env_name else ""
    f2p_cmd = f"{cmd_prefix}python -m pytest -v --tb=short -x {' '.join(f2p[:10])}"
    print(f"  F2P: {f2p_cmd[:150]}...")
    rc, out, err = run_shell(f"cd {repo_dir} && {f2p_cmd}", timeout=300)
    
    if rc == 0:
        # Run P2P regression tests
        p2p = instance.get("PASS_TO_PASS", [])
        if p2p:
            p2p_cmd = f"{cmd_prefix}python -m pytest -v --tb=short -x {' '.join(p2p[:10])}"
            rc2, out2, err2 = run_shell(f"cd {repo_dir} && {p2p_cmd}", timeout=300)
            if rc2 != 0:
                return {"resolved": False, "error": f"P2P regression: {(out2+err2)[:300]}"}
        return {"resolved": True, "test_output": (out + err)[:500]}
    
    return {"resolved": False, "error": f"F2P: {len(re.findall(r'FAILED', out+err))} failures", "test_output": (out + err)[:500]}


def main():
    print(f"🚀 CASCADE SMOKE TEST")
    print(f"   Instance: {INSTANCE_ID}")
    print(f"   T1: {T1_MODEL}")
    print(f"   T2: {T2_MODEL}")
    print(f"   Time: {datetime.now().isoformat()}")
    
    # Load instance
    print("\n[1/5] Loading SWE-bench_Verified...")
    ds = datasets.load_dataset("princeton-nlp/SWE-bench_Verified", split="test")
    instance = None
    for row in ds:
        if row["instance_id"] == INSTANCE_ID:
            instance = dict(row)
            break
    
    if not instance:
        print(f"❌ Instance {INSTANCE_ID} not found!")
        trackio.alert("Instance not found", f"{INSTANCE_ID} not in SWE-bench_Verified", level="ERROR")
        sys.exit(1)
    
    print(f"   Repo: {instance['repo']}")
    print(f"   Base: {instance['base_commit'][:12]}")
    print(f"   F2P tests: {len(instance.get('FAIL_TO_PASS', []))}")
    print(f"   P2P tests: {len(instance.get('PASS_TO_PASS', []))}")
    
    trackio.log({
        "event": "instance_loaded",
        "repo": instance["repo"],
        "f2p_count": len(instance.get("FAIL_TO_PASS", [])),
        "p2p_count": len(instance.get("PASS_TO_PASS", [])),
    })
    
    with tempfile.TemporaryDirectory(prefix="aco_smoke_") as tmpdir:
        repo_dir = Path(tmpdir) / "repo"
        env_name = f"aco_{INSTANCE_ID.replace('__','_').replace('-','_')[:30]}"
        
        # Clone
        print(f"\n[2/5] Cloning repo...")
        t0 = time.time()
        repo = instance["repo"]
        url = f"https://github.com/{repo}.git"
        rc, out, err = run_shell(f"git clone --depth 50 {url} {repo_dir}", timeout=300)
        if rc != 0:
            rc, out, err = run_shell(f"git clone {url} {repo_dir}", timeout=600)
        clone_time = time.time() - t0
        
        if rc != 0:
            print(f"❌ Clone failed: {err[:300]}")
            trackio.alert("Clone failed", err[:200], level="ERROR")
            sys.exit(1)
        print(f"   Done ({clone_time:.0f}s)")
        
        # Set up conda env
        print(f"\n[3/5] Setting up conda environment...")
        t0 = time.time()
        
        env_commit = instance.get("environment_setup_commit", "")
        if env_commit:
            run_shell(f"cd {repo_dir} && git fetch origin {env_commit}", timeout=60)
            run_shell(f"cd {repo_dir} && git checkout {env_commit}", timeout=30)
        
        env_yml = None
        for c in ["environment.yml", "dev/environment.yml", ".github/environment.yml",
                   "ci/environment.yml"]:
            if (repo_dir / c).exists():
                env_yml = c
                break
        
        if env_yml:
            print(f"   Found: {env_yml}")
            rc, out, err = run_shell(f"cd {repo_dir} && conda env create -f {env_yml} -n {env_name} --quiet", timeout=600)
        else:
            print(f"   No environment.yml, creating basic env")
            rc, out, err = run_shell(f"conda create -n {env_name} python=3.10 pip -y --quiet", timeout=300)
        
        if rc != 0:
            print(f"❌ Conda env creation failed: {err[:300]}")
            trackio.alert("Conda env failed", err[:300], level="ERROR")
            sys.exit(1)
        
        # Install repo at base_commit
        base_commit = instance["base_commit"]
        run_shell(f"cd {repo_dir} && git fetch origin {base_commit}", timeout=60)
        run_shell(f"cd {repo_dir} && git checkout {base_commit}", timeout=30)
        rc, out, err = run_shell(f"cd {repo_dir} && conda run -n {env_name} pip install -e . --quiet", timeout=300)
        if rc != 0:
            print(f"⚠️ pip install had issues: {err[:200]}")
        
        env_time = time.time() - t0
        print(f"   Done ({env_time:.0f}s)")
        
        # Run cascade agent
        print(f"\n[4/5] Running cascade agent (T1→T2)...")
        t0 = time.time()
        agent_result = run_cascade_agent(instance, repo_dir)
        agent_time = time.time() - t0
        
        if not agent_result["patch"]:
            print(f"❌ No patch produced")
            trackio.alert("No patch", "Cascade produced no patch", level="ERROR")
            sys.exit(1)
        
        print(f"\n   ✅ Patch: {len(agent_result['patch'])} chars")
        print(f"   Tier: {agent_result['tier']}")
        print(f"   Turns: {agent_result['turns']}")
        print(f"   Tokens: {agent_result['input_tokens']} in + {agent_result['output_tokens']} out")
        print(f"   Time: {agent_time:.0f}s")
        
        trackio.log({
            "event": "agent_done",
            "tier": agent_result["tier"],
            "turns": agent_result["turns"],
            "input_tokens": agent_result["input_tokens"],
            "output_tokens": agent_result["output_tokens"],
            "agent_time": agent_time,
        })
        
        # Verify
        print(f"\n[5/5] Verifying patch...")
        t0 = time.time()
        verify_result = verify_patch(instance, agent_result["patch"], repo_dir, env_name)
        verify_time = time.time() - t0
        
        trackio.log({
            "event": "verification_done",
            "resolved": verify_result["resolved"],
            "verify_time": verify_time,
        })
        
        # Final result
        print(f"\n{'='*60}")
        print(f"{'✅ RESOLVED' if verify_result['resolved'] else '❌ FAILED'}")
        print(f"{'='*60}")
        print(f"  Instance: {INSTANCE_ID}")
        print(f"  Tier: {agent_result['tier']}")
        print(f"  Turns: {agent_result['turns']}")
        print(f"  Time: clone={clone_time:.0f}s env={env_time:.0f}s agent={agent_time:.0f}s verify={verify_time:.0f}s")
        
        if not verify_result["resolved"]:
            print(f"  Error: {verify_result.get('error', 'unknown')[:300]}")
            trackio.alert("Not resolved", verify_result.get('error', 'unknown')[:300], level="WARN")
        else:
            trackio.alert("✅ Resolved!", f"Instance {INSTANCE_ID} resolved via {agent_result['tier']}", level="INFO")
        
        # Save final result
        final = {
            "instance_id": INSTANCE_ID,
            "repo": instance["repo"],
            "timestamp": datetime.now().isoformat(),
            "resolved": verify_result["resolved"],
            "tier": agent_result["tier"],
            "turns": agent_result["turns"],
            "input_tokens": agent_result["input_tokens"],
            "output_tokens": agent_result["output_tokens"],
            "clone_time": clone_time,
            "env_time": env_time,
            "agent_time": agent_time,
            "verify_time": verify_time,
            "patch_preview": agent_result["patch"][:500],
            "error": verify_result.get("error"),
        }
        
        with open("smoke_result.json", "w") as f:
            json.dump(final, f, indent=2)
        print(f"\nSaved: smoke_result.json")
        
        # Cleanup
        run_shell(f"conda env remove -n {env_name} -y --quiet", timeout=30)
    
    return 0 if verify_result["resolved"] else 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as e:
        print(f"💥 CRASH: {e}")
        traceback.print_exc()
        trackio.alert("Crash", str(e)[:300], level="ERROR")
        sys.exit(1)