"""
Execute open plot-skills beads tasks. 
Optimized for Beads: re-checks 'bd ready' per task and forces syncs to keep CLI/Agent in sync.
"""
import json
import subprocess
import os
import re
import sys

# Constants
VENV = "/home/russ/Documents/projects/Olympus/montera/scripts/env_montera"

def run_command(cmd, cwd=None):
    """Executes a command using the venv's environment and returns result."""
    env = os.environ.copy()
    # Add venv/bin to path for direct access to plotting tools
    env["PATH"] = os.path.join(VENV, "bin") + os.pathsep + env.get("PATH", "")
    
    return subprocess.run(
        cmd, 
        shell=True, 
        cwd=cwd or os.getcwd(), 
        env=env, 
        executable="/bin/bash"
    )

def bd_query_ready(cwd):
    """Fetches currently ready tasks as JSON."""
    r = subprocess.run(
        ["bd", "ready", "--json"],
        cwd=cwd,
        capture_output=True,
        text=True,
    )
    if r.returncode != 0:
        return []
    try:
        data = json.loads(r.stdout)
        # Filter for only plot-skills related tasks if necessary
        return [t for t in data if "plot-skills" in t.get("labels", []) or "plot" in t.get("title", "").lower()]
    except json.JSONDecodeError:
        return []

def date_dirs(base):
    """Yield (base_path, date_path) where date_path has gnssdata."""
    gnss = os.path.join(base, "gnssdata")
    if os.path.isdir(gnss):
        yield base, base
        return
    try:
        # Sort to ensure predictable execution order
        for name in sorted(os.listdir(base)):
            if re.match(r"^[0-9]{8}$", name):
                sub = os.path.join(base, name)
                if os.path.isdir(os.path.join(sub, "gnssdata")):
                    yield base, sub
    except OSError:
        pass

def execute_one(t, repo):
    """Handles the logic for a single bead."""
    title = t.get("title", "")
    tid = t["id"]
    
    # 1. Claim the task immediately
    print(f"--- Processing {tid} ---")
    run_command(f"bd update {tid} --status in_progress", cwd=repo)
    
    # 2. Parse the command requirements
    m = re.search(r"Run plot_(snr_diff_hdf5|gnss_hdf5): ([^\s(]+)", title)
    if not m:
        print(f"SKIP {tid}: Title format mismatch.")
        run_command(f"bd update {tid} --status closed --notes 'Skipped: Format mismatch'", cwd=repo)
        return

    tool, dirname = m.group(1), m.group(2).strip()
    if not os.path.isdir(dirname):
        print(f"SKIP {tid}: Directory not found: {dirname}")
        run_command(f"bd update {tid} --status closed --notes 'Skipped: Dir not found'", cwd=repo)
        return

    # 3. Execution Loop
    success = True
    if tool == "snr_diff_hdf5":
        for _, date_dir in date_dirs(dirname):
            cmd = f"plot_snr_diff_hdf5 {os.path.abspath(date_dir)} BX992 SUMO2"
            print(f"EXEC: {cmd}")
            r = run_command(cmd, cwd=repo)
            if r.returncode != 0:
                success = False
    else:
        dut = "BX992" if "BX992" in title else "SUMO2"
        for _, date_dir in date_dirs(dirname):
            rec_dir = os.path.join(date_dir, "gnssdata", dut)
            if not os.path.isdir(rec_dir):
                continue
            cmd = f"plot_gnss_hdf5 {os.path.abspath(rec_dir)}"
            print(f"EXEC: {cmd}")
            r = run_command(cmd, cwd=repo)
            if r.returncode != 0:
                success = False

    # 4. Finalize the bead
    status = "closed" if success else "open"
    notes = "Completed successfully via script." if success else "One or more plots failed."
    
    run_command(f"bd update {tid} --status {status} --notes '{notes}'", cwd=repo)
    
    # 5. CRITICAL: Sync after every task to update the JSONL files for Git and Agent visibility
    run_command("bd sync", cwd=repo)
    print(f"DONE {tid}: {status}")

def main():
    repo = os.getcwd()
    
    print("Beads Task Runner Started.")
    
    while True:
        # We query 'ready' every iteration. This ensures that if a previous 
        # task's closure unblocked new tasks, we see them immediately.
        ready_tasks = bd_query_ready(repo)
        
        if not ready_tasks:
            print("No more ready tasks found. Exiting.")
            break
            
        # Execute only the first ready task, then loop back to re-query.
        # This prevents the script from working on stale data if the 
        # dependency graph changes during execution.
        target_task = ready_tasks[0]
        execute_one(target_task, repo)

if __name__ == "__main__":
    main()