"""
Cascade Smoke Test — Self-Contained HF Job Script

Tests the cascade agent end-to-end on one SWE-bench instance.
Installs miniconda if needed, sets up environment, runs agent,
verifies patch.

Usage via hf_jobs:
  operation: run
  script: "https://huggingface.co/narcolepticchicken/agent-cost-optimizer/resolve/main/smoke_test.py"
  dependencies: ["huggingface_hub", "datasets", "trackio"]
  hardware: a10g-largex2
  timeout: 4h
  env:
    INSTANCE_ID: "django__django-14315"
"""

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

# ============================================================
# Bootstrap: ensure conda is available
# ============================================================

def ensure_conda():
    """Install miniconda if not present."""
    for path in [os.path.expanduser("~/miniconda3/bin/conda"), "/opt/conda/bin/conda"]:
        if os.path.exists(path):
            return path
    
    print("📦 Installing Miniconda...")
    rc, out, err = subprocess_run(
        "wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh && "
        "bash /tmp/miniconda.sh -b -p $HOME/miniconda3",
        timeout=120
    )
    if rc != 0:
        print(f"Conda install failed: {err}")
        return None
    
    conda_path = os.path.expanduser("~/miniconda3/bin/conda")
    # Add to PATH for this session
    os.environ["PATH"] = os.path.expanduser("~/miniconda3/bin:") + os.environ.get("PATH", "")
    return conda_path


def subprocess_run(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 conda_run(conda, subcmd, timeout=300):
    """Run a conda command with timeout."""
    return subprocess_run(f"{conda} {subcmd}", timeout=timeout)


# ============================================================
# Agent
# ============================================================

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(instance, repo_dir, conda, env_name):
    """Run T1 → T2 cascade agent."""
    from huggingface_hub import InferenceClient
    
    T1 = "meta-llama/Llama-3.1-8B-Instruct"
    T2 = "meta-llama/Llama-3.3-70B-Instruct"
    
    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. 
Use pytest to verify. Be thorough but efficient."""

    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": f"PROBLEM:\n{problem}\n\nStart by exploring the repository structure."}
    ]
    
    tiers = [("T1", T1, 30), ("T2", T2, 30)]
    
    for tier_name, model_id, max_turns in tiers:
        print(f"\n[{tier_name}] {model_id}")
        client = 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} tok, {len(text)} ch")
            
            patch = extract_patch(text)
            if patch:
                print(f"  ✅ PATCH ({len(patch)} ch)")
                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()
                cmd = cmd.replace("pytest", f"{conda} run -n {env_name} python -m pytest") if "pytest" in cmd else cmd
                print(f"  $ {cmd[:120]}")
                rc, stdout, stderr = subprocess_run(cmd, cwd=str(repo_dir), timeout=60)
                output = (stdout + stderr)[:1500]
                if rc != 0:
                    output += f" [EXIT:{rc}]"
                messages.append({"role": "user", "content": f"<output>\n{output}\n</output>"})
            
            if "<submit>" in text:
                break
    
    return {"patch": None, "tier": None, "turns": 0, "input_tokens": 0, "output_tokens": 0}


def verify_patch(instance, model_patch, repo_dir, conda, env_name):
    """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
    subprocess_run(f"cd {repo_dir} && git checkout -f {base_commit}", timeout=30)
    subprocess_run(f"cd {repo_dir} && git clean -fd", timeout=30)
    
    # Apply model patch
    (Path(repo_dir) / "_aco.patch").write_text(model_patch)
    rc, out, err = subprocess_run(f"cd {repo_dir} && git apply --check _aco.patch", timeout=10)
    if rc != 0:
        return {"resolved": False, "error": f"patch --check: {err[:200]}"}
    rc, out, err = subprocess_run(f"cd {repo_dir} && git apply _aco.patch", timeout=10)
    if rc != 0:
        return {"resolved": False, "error": f"patch apply: {err[:200]}"}
    
    # Apply test_patch
    (Path(repo_dir) / "_aco_test.patch").write_text(test_patch)
    subprocess_run(f"cd {repo_dir} && git apply --check _aco_test.patch && git apply _aco_test.patch || git apply --reject _aco_test.patch", timeout=10)
    
    # Run F2P tests
    f2p_str = ' '.join(f2p[:10])
    cmd = f"cd {repo_dir} && {conda} run -n {env_name} python -m pytest -v --tb=short -x {f2p_str}"
    print(f"  F2P: pytest {' '.join(f2p[:3])}...")
    rc, out, err = subprocess_run(cmd, timeout=300)
    
    if rc == 0:
        p2p = instance.get("PASS_TO_PASS", [])
        if p2p:
            p2p_str = ' '.join(p2p[:10])
            cmd2 = f"cd {repo_dir} && {conda} run -n {env_name} python -m pytest -v --tb=short -x {p2p_str}"
            rc2, out2, err2 = subprocess_run(cmd2, timeout=300)
            if rc2 != 0:
                return {"resolved": False, "error": f"P2P: {(out2+err2)[:200]}"}
        return {"resolved": True, "test_output": (out + err)[:500]}
    
    failures = len(re.findall(r'FAILED', out + err))
    return {"resolved": False, "error": f"{failures} F2P failures", "test_output": (out + err)[:500]}


def setup_environment(conda, instance, repo_dir, env_name):
    """Create conda environment for the repo."""
    env_commit = instance.get("environment_setup_commit", "")
    
    if env_commit:
        subprocess_run(f"cd {repo_dir} && git fetch origin {env_commit}", timeout=60)
        subprocess_run(f"cd {repo_dir} && git checkout {env_commit}", timeout=30)
    
    # Find environment.yml
    env_yml = None
    for c in ["environment.yml", "dev/environment.yml", ".github/environment.yml",
               "ci/environment.yml"]:
        if (Path(repo_dir) / c).exists():
            env_yml = c
            break
    
    if env_yml:
        print(f"  Using {env_yml}")
        rc, out, err = conda_run(conda, f"env create -f {repo_dir}/{env_yml} -n {env_name} --quiet", timeout=600)
    else:
        print(f"  Creating basic python=3.10 env")
        rc, out, err = conda_run(conda, f"create -n {env_name} python=3.10 pip -y --quiet", timeout=300)
    
    if rc != 0:
        print(f"  ⚠️ env creation failed: {err[:200]}")
        # Try with just pip
        rc2, out2, err2 = conda_run(conda, f"create -n {env_name} python=3.10 pip -y --quiet", timeout=300)
        if rc2 != 0:
            return False, f"conda env: {err[:200]}"
    
    # Install repo
    base_commit = instance["base_commit"]
    subprocess_run(f"cd {repo_dir} && git fetch origin {base_commit}", timeout=60)
    subprocess_run(f"cd {repo_dir} && git checkout {base_commit}", timeout=30)
    rc, out, err = conda_run(conda, f"run -n {env_name} pip install -e . --quiet", 
                             cwd=str(repo_dir), timeout=300)
    if rc != 0:
        print(f"  ⚠️ pip install: {err[:200]}")
        # Try without -e
        rc2, out2, err2 = conda_run(conda, f"run -n {env_name} pip install . --quiet",
                                     cwd=str(repo_dir), timeout=300)
    
    return True, ""


def main():
    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"
    
    print("="*60)
    print(f"🚀 CASCADE SMOKE TEST — {datetime.now().isoformat()}")
    print("="*60)
    print(f"  Instance: {INSTANCE_ID}")
    print(f"  T1: {T1_MODEL}")
    print(f"  T2: {T2_MODEL}")
    
    # Ensure conda
    conda = ensure_conda()
    if not conda:
        print("❌ Cannot install conda")
        sys.exit(1)
    print(f"\n✅ Conda: {conda}")
    
    # Load instance
    print(f"\n[1/5] Loading instance...")
    import datasets
    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_ID} not found!")
        sys.exit(1)
    
    print(f"  Repo: {instance['repo']}")
    print(f"  Base: {instance['base_commit'][:12]}")
    print(f"  F2P: {len(instance.get('FAIL_TO_PASS', []))} tests")
    print(f"  P2P: {len(instance.get('PASS_TO_PASS', []))} tests")
    
    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 {instance['repo']}...")
        t0 = time.time()
        url = f"https://github.com/{instance['repo']}.git"
        rc, out, err = subprocess_run(f"git clone --depth 50 {url} {repo_dir}", timeout=300)
        if rc != 0:
            rc, out, err = subprocess_run(f"git clone {url} {repo_dir}", timeout=600)
        clone_t = time.time() - t0
        if rc != 0:
            print(f"❌ Clone: {err[:200]}")
            sys.exit(1)
        print(f"  Done ({clone_t:.0f}s)")
        
        # Environment
        print(f"\n[3/5] Setting up conda env '{env_name}'...")
        t0 = time.time()
        ok, err = setup_environment(conda, instance, repo_dir, env_name)
        env_t = time.time() - t0
        if not ok:
            print(f"❌ Env setup: {err}")
            sys.exit(1)
        print(f"  Done ({env_t:.0f}s)")
        
        # Agent
        print(f"\n[4/5] Running cascade agent...")
        t0 = time.time()
        agent = run_cascade(instance, repo_dir, conda, env_name)
        agent_t = time.time() - t0
        
        if not agent["patch"]:
            print(f"❌ No patch!")
            sys.exit(1)
        
        print(f"  ✅ {agent['tier']}, {agent['turns']} turns, {agent['input_tokens']}+{agent['output_tokens']} tokens")
        
        # Verify
        print(f"\n[5/5] Verifying...")
        t0 = time.time()
        verify = verify_patch(instance, agent["patch"], repo_dir, conda, env_name)
        verify_t = time.time() - t0
        
        # Result
        status = "✅ RESOLVED" if verify["resolved"] else "❌ FAILED"
        print(f"\n{'='*60}")
        print(f"{status}")
        print(f"{'='*60}")
        print(f"  Times: clone={clone_t:.0f}s env={env_t:.0f}s agent={agent_t:.0f}s verify={verify_t:.0f}s")
        if not verify["resolved"]:
            print(f"  Error: {verify.get('error', 'unknown')[:300]}")
        
        # Save
        result = {
            "instance_id": INSTANCE_ID,
            "repo": instance["repo"],
            "resolved": verify["resolved"],
            "tier": agent["tier"],
            "turns": agent["turns"],
            "input_tokens": agent["input_tokens"],
            "output_tokens": agent["output_tokens"],
            "times": {"clone": clone_t, "env": env_t, "agent": agent_t, "verify": verify_t},
            "error": verify.get("error"),
            "patch_preview": agent["patch"][:500],
        }
        
        with open("smoke_result.json", "w") as f:
            json.dump(result, f, indent=2)
        print(f"\nSaved: smoke_result.json")
        
        # Cleanup
        conda_run(conda, f"env remove -n {env_name} -y --quiet", timeout=30)
    
    return 0 if verify["resolved"] else 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as e:
        print(f"💥 {e}")
        traceback.print_exc()
        sys.exit(1)