"""Researcher's verification agent — verify contributor branches in isolated worktrees."""
from __future__ import annotations
import os
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from .client import NbenchClient
from .config import load_config


def _parse_last_result(results_file, config):
    """Parse the last line of the results file to get the metric value."""
    if not os.path.exists(results_file):
        return None
    with open(results_file) as f:
        lines = [l.strip() for l in f.readlines() if l.strip()]
    if not lines:
        return None
    # TSV format: commit_hash\tmetric_value\tstatus\tdescription
    parts = lines[-1].split("\t")
    if len(parts) >= 2:
        try:
            return float(parts[1])
        except ValueError:
            return None
    return None


def _verify_branch(client, project_id, branch, config, repo_path):
    """Verify a single branch in an isolated git worktree."""
    branch_name = branch["branch_name"]
    branch_id = branch["id"]
    safe_name = branch_name.replace("/", "_")
    worktree_path = os.path.join(repo_path, ".worktrees", safe_name)

    print(f"[watch] === Verifying: {branch_name} ===")
    contributor = branch.get("contributor_handle", "unknown")
    print(f"[watch] Contributor: @{contributor}")

    try:
        # Mark as verifying
        client.update_branch(project_id, branch_id, status="verifying")

        # Fetch the branch
        result = subprocess.run(
            ["git", "fetch", "origin", branch_name],
            cwd=repo_path, capture_output=True, text=True
        )
        if result.returncode != 0:
            client.update_branch(project_id, branch_id,
                status="rejected", error_log=f"Failed to fetch branch: {result.stderr[:500]}")
            print(f"[watch] REJECTED: could not fetch branch")
            return

        # Check branch is not behind main
        result = subprocess.run(
            ["git", "log", f"origin/main..origin/{branch_name}", "--oneline"],
            cwd=repo_path, capture_output=True, text=True
        )
        if not result.stdout.strip():
            client.update_branch(project_id, branch_id,
                status="rejected", error_log="Branch has no new commits ahead of main")
            print(f"[watch] REJECTED: no commits ahead of main")
            return

        # Create worktree
        os.makedirs(os.path.dirname(worktree_path), exist_ok=True)
        subprocess.run(
            ["git", "worktree", "add", worktree_path, f"origin/{branch_name}"],
            cwd=repo_path, capture_output=True, check=True
        )

        # Clean any existing results
        results_path = os.path.join(worktree_path, config.get("results", "results.tsv"))
        if os.path.exists(results_path):
            os.remove(results_path)

        # Check diff for suspicious changes
        diff_result = subprocess.run(
            ["git", "diff", "origin/main", f"origin/{branch_name}", "--stat"],
            cwd=repo_path, capture_output=True, text=True
        )
        diff_files = [l.split("|")[0].strip() for l in diff_result.stdout.strip().split("\n") if "|" in l]
        suspicious = [f for f in diff_files if any(s in f.lower() for s in [
            ".env", ".nbench", "credentials", ".ssh", ".git/", "dockerfile",
            "makefile", ".github", "package.json", "go.mod", "requirements.txt"
        ])]
        if suspicious:
            client.update_branch(project_id, branch_id,
                status="rejected", error_log=f"Suspicious files modified: {', '.join(suspicious)}")
            print(f"[watch] REJECTED: suspicious files: {suspicious}")
            return

        # Run training in a clean environment (strip secrets)
        train_cmd = config.get("train", "python train.py")
        clean_env = {k: v for k, v in os.environ.items()
                    if not any(s in k.upper() for s in ["KEY", "SECRET", "TOKEN", "PASSWORD", "NB_"])}
        clean_env["PATH"] = os.environ.get("PATH", "")
        clean_env["HOME"] = os.environ.get("HOME", "")
        clean_env["NB_RESULTS_FILE"] = config.get("results", "results.tsv")

        print(f"[watch] Running: {train_cmd} (sandboxed env)")
        result = subprocess.run(
            train_cmd, shell=True, cwd=worktree_path,
            capture_output=True, text=True, timeout=3600,
            env=clean_env
        )

        if result.returncode != 0:
            error = result.stderr[:1000] if result.stderr else "Training exited with non-zero code"
            client.update_branch(project_id, branch_id,
                status="rejected", error_log=error)
            print(f"[watch] REJECTED: training failed")
            if result.stderr:
                print(f"[watch] stderr: {result.stderr[:200]}")
            return

        # Parse result
        verified_metric = _parse_last_result(results_path, config)
        if verified_metric is None:
            client.update_branch(project_id, branch_id,
                status="rejected", error_log="No results produced or could not parse metric")
            print(f"[watch] REJECTED: no results")
            return

        # Compare with current best
        project = client.get_project(project_id)
        current_best = project.get("current_best_val_bpb", float("inf"))
        direction = config.get("direction", "minimize")

        improved = (
            (direction == "minimize" and verified_metric < current_best) or
            (direction == "maximize" and verified_metric > current_best)
        )

        if improved:
            # Merge
            subprocess.run(["git", "checkout", "main"], cwd=repo_path, capture_output=True)
            merge_result = subprocess.run(
                ["git", "merge", f"origin/{branch_name}", "--no-edit"],
                cwd=repo_path, capture_output=True, text=True
            )

            if merge_result.returncode != 0:
                # Merge conflict
                subprocess.run(["git", "merge", "--abort"], cwd=repo_path, capture_output=True)
                client.update_branch(project_id, branch_id,
                    status="conflict", verified_metric=verified_metric,
                    error_log="Merge conflict — rebase on main and try again")
                print(f"[watch] CONFLICT: verified improvement ({verified_metric}) but merge conflict")
                return

            subprocess.run(["git", "push", "origin", "main"], cwd=repo_path, capture_output=True)

            client.update_branch(project_id, branch_id,
                status="merged", verified_metric=verified_metric)
            print(f"[watch] MERGED: {current_best} -> {verified_metric}")
        else:
            client.update_branch(project_id, branch_id,
                status="rejected", verified_metric=verified_metric,
                error_log=f"No improvement: {verified_metric} vs best {current_best} ({direction})")
            print(f"[watch] REJECTED: no improvement ({verified_metric} vs {current_best})")

    except subprocess.TimeoutExpired:
        client.update_branch(project_id, branch_id,
            status="rejected", error_log="Training timed out (1 hour)")
        print(f"[watch] REJECTED: timeout")
    except Exception as e:
        client.update_branch(project_id, branch_id,
            status="rejected", error_log=str(e)[:1000])
        print(f"[watch] ERROR: {e}")
    finally:
        # Clean up worktree
        subprocess.run(
            ["git", "worktree", "remove", worktree_path, "--force"],
            cwd=repo_path, capture_output=True
        )


def run_watch(api_url, api_key, project_id, repo_path=".", max_concurrent=3, interval=30):
    """Main watch loop: poll for pending branches, verify concurrently."""
    client = NbenchClient(api_url=api_url, api_key=api_key)
    config = load_config(repo_path)

    print(f"[watch] Researcher's verification agent")
    print(f"[watch] Project:  {project_id}")
    print(f"[watch] Repo:     {os.path.abspath(repo_path)}")
    print(f"[watch] Metric:   {config['metric']} ({config['direction']})")
    print(f"[watch] Train:    {config['train']}")
    print(f"[watch] Workers:  {max_concurrent}")
    print()
    print(f"[watch] Security: branches run in isolated worktrees with stripped env vars.")
    print(f"[watch] Branches modifying config/dependency files are auto-rejected.")
    print()

    active_branches = set()

    with ThreadPoolExecutor(max_workers=max_concurrent) as pool:
        while True:
            try:
                branches = client.get_branches(project_id)
                pending = [
                    b for b in branches
                    if b["status"] == "pending" and b["branch_name"] not in active_branches
                ]

                for branch in pending:
                    name = branch["branch_name"]
                    active_branches.add(name)
                    print(f"[watch] Queuing verification: {name}")
                    future = pool.submit(
                        _verify_branch, client, project_id, branch, config, repo_path
                    )
                    future.add_done_callback(lambda f, n=name: active_branches.discard(n))

                if not pending and not active_branches:
                    print(f"[watch] No pending branches. Checking in {interval}s...")

                time.sleep(interval)

            except KeyboardInterrupt:
                print(f"\n[watch] Stopped.")
                break
            except Exception as e:
                print(f"[watch] Error: {e}")
                time.sleep(interval)
