#!/usr/bin/env python3.12
"""Find next batch or sites from registry based on status.

Usage:
    find-next.py --registry PATH --mode batch|sites [--phase design|implement]
    find-next.py --registry PATH --mode batch --phase implement --claim --lock-file PATH

Modes:
    batch   - Find next batch ID to work on
    sites   - Find sites that need processing
    status  - Show status counts
    active  - Find currently active batch

Claim mode (--claim):
    Atomically finds next work AND updates status to claim it.
    Uses fcntl locking to prevent race conditions between subagents.
    - Design: B → d (claim for design)
    - Implement: D → O or O → i (claim for implement)

Exit codes:
    0 - Work found successfully
    1 - No work available (expected condition, not an error)
    2 - Actual error (permission denied, OS error, etc.)

Output (JSON):
    Success: {"batch": "003", "sites": ["domain1", "domain2"], "exit_code": 0, "claimed": true}
    No work: {"batch": null, "sites": [], "reason": "no_work_found", "exit_code": 1}
    Error:   {"error": "permission_denied", "message": "...", "exit_code": 2}

Output (simple):
    Success: "003" or "domain1,domain2"
    No work: "ERROR:no_work_found"
    Error:   "ERROR:permission_denied:details..."
"""
import argparse
import datetime
import fcntl
import json
import pathlib
import sys
from collections import defaultdict
from typing import Dict, List, Optional


def parse_registry(registry_path: pathlib.Path) -> tuple[List[Dict[str, str]], List[str]]:
    """Parse REGISTRY.md and return list of site entries and raw lines."""
    if not registry_path.exists():
        return [], []

    entries = []
    lines = registry_path.read_text().splitlines()
    for line in lines:
        if line.startswith("|") and "Domain" not in line and "--------" not in line:
            cols = [c.strip() for c in line.strip("|").split("|")]
            if len(cols) >= 6:
                entries.append({
                    "domain": cols[0],
                    "title": cols[1],
                    "description": cols[2],
                    "status": cols[3],
                    "batch": cols[4],
                    "updated": cols[5],
                })
    return entries, lines


def update_registry_status(
    registry_path: pathlib.Path,
    domains: List[str],
    new_status: str,
) -> None:
    """Update status for domains in registry (caller must hold lock)."""
    if not registry_path.exists():
        return

    lines = registry_path.read_text().splitlines()
    new_lines = []
    now = datetime.datetime.now().isoformat(timespec="seconds")

    for line in lines:
        if line.startswith("|") and "Domain" not in line and "--------" not in line:
            cols = [c.strip() for c in line.strip("|").split("|")]
            if len(cols) >= 6 and cols[0] in domains:
                cols[3] = new_status
                cols[5] = now
                line = f"| {cols[0]} | {cols[1]} | {cols[2]} | {cols[3]} | {cols[4]} | {cols[5]} |"
        new_lines.append(line)

    registry_path.write_text("\n".join(new_lines) + "\n")


def find_next_design_batch(entries: List[Dict[str, str]]) -> Optional[Dict]:
    """Find next batch for design phase.

    Priority:
    1. Batch with 'd' status (in-progress) - continue
    2. Batch with 'B' status (batched, ready to start)
    3. Sites with '-' status (need batching first)
    """
    by_batch: Dict[str, List[Dict[str, str]]] = defaultdict(list)
    unassigned: List[Dict[str, str]] = []

    for e in entries:
        if e["status"] == "-":
            unassigned.append(e)
        elif e["batch"]:
            by_batch[e["batch"]].append(e)

    # Priority 1: In-progress design (d)
    for batch_id, sites in sorted(by_batch.items()):
        d_sites = [s for s in sites if s["status"] == "d"]
        if d_sites:
            return {
                "batch": batch_id,
                "sites": [s["domain"] for s in d_sites],
                "status": "d",
                "count": len(d_sites),
                "action": "continue-design",
                "claim_status": None,  # Already claimed
            }

    # Priority 2: Ready for design (B)
    for batch_id, sites in sorted(by_batch.items()):
        b_sites = [s for s in sites if s["status"] == "B"]
        if b_sites:
            return {
                "batch": batch_id,
                "sites": [s["domain"] for s in b_sites],
                "status": "B",
                "count": len(b_sites),
                "action": "start-design",
                "claim_status": "d",  # Claim by setting to d
            }

    # Priority 3: Unassigned sites need batching
    if unassigned:
        return {
            "batch": None,
            "sites": [s["domain"] for s in unassigned],
            "status": "-",
            "count": len(unassigned),
            "action": "create-batch",
            "claim_status": None,  # Cannot claim unassigned
        }

    return None


def find_next_implement_batch(entries: List[Dict[str, str]]) -> Optional[Dict]:
    """Find next batch for implement phase.

    Priority:
    1. Batch with 'i' status (in-progress) - continue
    2. Batch with 'O' status (ready to start)
    3. Batch with 'D' status (design done, needs O transition)
    """
    by_batch: Dict[str, List[Dict[str, str]]] = defaultdict(list)

    for e in entries:
        if e["batch"]:
            by_batch[e["batch"]].append(e)

    # Priority 1: In-progress implement (i)
    for batch_id, sites in sorted(by_batch.items()):
        i_sites = [s for s in sites if s["status"] == "i"]
        if i_sites:
            return {
                "batch": batch_id,
                "sites": [s["domain"] for s in i_sites],
                "status": "i",
                "count": len(i_sites),
                "action": "continue-implement",
                "claim_status": None,  # Already claimed
            }

    # Priority 2: Ready for implement (O)
    for batch_id, sites in sorted(by_batch.items()):
        o_sites = [s for s in sites if s["status"] == "O"]
        if o_sites:
            return {
                "batch": batch_id,
                "sites": [s["domain"] for s in o_sites],
                "status": "O",
                "count": len(o_sites),
                "action": "start-implement",
                "claim_status": "i",  # Claim by setting to i
            }

    # Priority 3: Design done, needs transition (D)
    for batch_id, sites in sorted(by_batch.items()):
        d_sites = [s for s in sites if s["status"] == "D"]
        if d_sites:
            return {
                "batch": batch_id,
                "sites": [s["domain"] for s in d_sites],
                "status": "D",
                "count": len(d_sites),
                "action": "transition-to-implement",
                "claim_status": "O",  # First transition to O
            }

    return None


def find_active_batch(entries: List[Dict[str, str]]) -> Optional[Dict]:
    """Find batch with active work (d, i, or most recent non-Q)."""
    by_batch: Dict[str, List[Dict[str, str]]] = defaultdict(list)

    for e in entries:
        if e["batch"]:
            by_batch[e["batch"]].append(e)

    for batch_id, sites in sorted(by_batch.items(), reverse=True):
        active_statuses = {"d", "i", "B", "D", "O", "I"}
        active = [s for s in sites if s["status"] in active_statuses]
        if active:
            status_counts = defaultdict(int)
            for s in sites:
                status_counts[s["status"]] += 1
            return {
                "batch": batch_id,
                "sites": [s["domain"] for s in sites],
                "status_counts": dict(status_counts),
                "count": len(sites),
            }

    return None


def get_status_summary(entries: List[Dict[str, str]]) -> Dict[str, int]:
    """Get count of sites by status."""
    counts: Dict[str, int] = defaultdict(int)
    for e in entries:
        counts[e["status"]] += 1
    return dict(counts)


def main() -> int:
    parser = argparse.ArgumentParser(description="Find next batch or sites from registry")
    parser.add_argument("--registry", required=True, help="Path to REGISTRY.md")
    parser.add_argument("--mode", choices=["batch", "sites", "status", "active"],
                        default="batch", help="What to find")
    parser.add_argument("--phase", choices=["design", "implement"],
                        default="implement", help="Phase to find work for")
    parser.add_argument("--format", choices=["json", "simple"], default="json",
                        help="Output format")
    parser.add_argument("--claim", action="store_true",
                        help="Atomically claim the found batch/sites by updating status")
    parser.add_argument("--lock-file", default=None,
                        help="Lock file path (required with --claim, defaults to REGISTRY.lock)")
    args = parser.parse_args()

    registry = pathlib.Path(args.registry)

    # Determine lock file path
    if args.claim:
        if args.lock_file:
            lock_path = pathlib.Path(args.lock_file)
        else:
            lock_path = registry.parent / "REGISTRY.lock"
        lock_path.parent.mkdir(parents=True, exist_ok=True)
    else:
        lock_path = None

    # Exit codes:
    #   0 = work found successfully
    #   1 = no work available (expected, not an error)
    #   2 = actual error (registry issue, locking failure, etc.)

    # Execute with or without locking
    try:
        if args.claim and lock_path:
            # Atomic find + claim with fcntl lock
            with lock_path.open("w") as lf:
                fcntl.flock(lf, fcntl.LOCK_EX)
                try:
                    result = _find_and_claim(registry, args)
                finally:
                    fcntl.flock(lf, fcntl.LOCK_UN)
        else:
            # Read-only find (no locking needed for reads)
            result = _find_only(registry, args)
    except PermissionError as e:
        result = {"error": "permission_denied", "message": str(e), "exit_code": 2}
    except OSError as e:
        result = {"error": "os_error", "message": str(e), "exit_code": 2}
    except Exception as e:
        result = {"error": "unexpected", "message": str(e), "exit_code": 2}

    # Determine exit code
    if result is None:
        result = {"batch": None, "sites": [], "action": "none", "reason": "no_work_found", "exit_code": 1}
        exit_code = 1
    elif result.get("error"):
        exit_code = result.get("exit_code", 2)
    elif result.get("batch") or result.get("sites"):
        result["exit_code"] = 0
        exit_code = 0
    else:
        result["reason"] = "no_work_found"
        result["exit_code"] = 1
        exit_code = 1

    # Output result to stdout (always JSON for machine parsing, simple for human)
    if args.format == "simple":
        if exit_code == 0:
            if result.get("batch"):
                print(result["batch"])
            elif result.get("sites"):
                print(",".join(result["sites"]))
        else:
            # Print reason to stdout so caller can see it
            reason = result.get("error") or result.get("reason") or "unknown"
            message = result.get("message", "")
            print(f"ERROR:{reason}:{message}" if message else f"ERROR:{reason}")
    else:
        print(json.dumps(result, indent=2))

    return exit_code


def _find_only(registry: pathlib.Path, args) -> Optional[Dict]:
    """Find without claiming (read-only)."""
    entries, _ = parse_registry(registry)

    if not entries:
        return {"error": "no entries", "batch": None, "sites": []}

    if args.mode == "status":
        return {"status_counts": get_status_summary(entries), "total": len(entries)}
    elif args.mode == "active":
        return find_active_batch(entries)
    elif args.mode in ("batch", "sites"):
        if args.phase == "design":
            return find_next_design_batch(entries)
        else:
            return find_next_implement_batch(entries)

    return None


def _find_and_claim(registry: pathlib.Path, args) -> Optional[Dict]:
    """Find and atomically claim (caller must hold lock)."""
    entries, _ = parse_registry(registry)

    if not entries:
        return {"error": "no entries", "batch": None, "sites": [], "claimed": False}

    # Find next work
    if args.phase == "design":
        result = find_next_design_batch(entries)
    else:
        result = find_next_implement_batch(entries)

    if result is None:
        return {"batch": None, "sites": [], "action": "none", "message": "no work found", "claimed": False}

    # Claim if possible
    claim_status = result.get("claim_status")
    if claim_status and result.get("sites"):
        update_registry_status(registry, result["sites"], claim_status)
        result["claimed"] = True
        result["new_status"] = claim_status
        print(f"[find-next] Claimed {len(result['sites'])} sites with status {claim_status}", file=sys.stderr)
    else:
        result["claimed"] = False
        if result.get("action") in ("continue-design", "continue-implement"):
            print(f"[find-next] Continuing existing work on batch {result.get('batch')}", file=sys.stderr)

    return result


if __name__ == "__main__":
    raise SystemExit(main())