import argparse
import sys
from typing import Any


def audit(parts: list[dict[str, Any]], child_tool_counts: dict[str, int], worktree: str | None) -> list[dict[str, Any]]:
    """
    Pure core function to audit dispatch calls.
    
    Flags:
    - DEAD: No category, no subagent_type, AND no task_id.
    - BANNED: subagent_type starts with 'oh-my-claudecode:'.
    - NOWT: worktree is provided, but the prompt lacks a 'WORKTREE:' line.
    - IDLE: Child session exists but has 0 tool calls.
    """
    results = []
    
    for part in parts:
        if part.get("type") != "tool":
            continue
            
        tool_name = part.get("tool")
        if tool_name not in ("task", "call_omo_agent"):
            continue
            
        state = part.get("state", {})
        input_data = state.get("input", {})
        metadata = state.get("metadata", {})
        
        prompt = input_data.get("prompt", "")
        category = input_data.get("category")
        subagent_type = input_data.get("subagent_type")
        task_id = input_data.get("task_id")
        
        child_session_id = metadata.get("sessionId")
        
        flags = []
        
        # DEAD: no category, no subagent_type, no task_id
        if not category and not subagent_type and not task_id:
            flags.append("DEAD")
            
        # BANNED: subagent_type starts with 'oh-my-claudecode:'
        if subagent_type and subagent_type.startswith("oh-my-claudecode:"):
            flags.append("BANNED")
            
        # NOWT: worktree provided but prompt missing WORKTREE line
        if worktree and "WORKTREE:" not in prompt:
            flags.append("NOWT")
            
        # IDLE: child session exists but 0 tool calls
        if child_session_id and child_session_id in child_tool_counts and child_tool_counts[child_session_id] == 0:
            flags.append("IDLE")
            
        results.append({
            "tool": tool_name,
            "flags": flags,
            "sessionId": metadata.get("sessionId", "N/A")
        })
        
    return results

def main():
    parser = argparse.ArgumentParser(description="Audit OpenCode dispatch calls for guardrail violations.")
    parser.add_argument("session_id", help="The session ID to audit")
    parser.add_argument("--url", default="http://127.0.0.1:4097", help="OpenCode API URL")
    parser.add_argument("--worktree", help="The expected worktree path")
    
    args = parser.parse_args()
    
    import requests
    
    try:
        resp = requests.get(f"{args.url}/sessions/{args.session_id}")
        resp.raise_for_status()
        session_data = resp.json()
        parts = session_data.get("parts", [])
        
        child_tool_counts = {}
        child_sessions = [p.get("state", {}).get("metadata", {}).get("sessionId") 
                          for p in parts if p.get("type") == "tool" and "sessionId" in p.get("state", {}).get("metadata", {})]
        
        for csid in set(filter(None, child_sessions)):
            try:
                cs_resp = requests.get(f"{args.url}/sessions/{csid}")
                cs_resp.raise_for_status()
                cs_parts = cs_resp.json().get("parts", [])
                count = sum(1 for p in cs_parts if p.get("type") == "tool")
                child_tool_counts[csid] = count
            except Exception: # noqa: BLE001 CLI error handling: broad catch on HTTP requests
                child_tool_counts[csid] = -1
                
        findings = audit(parts, child_tool_counts, args.worktree)
        
        any_flagged = False
        for f in findings:
            if f["flags"]:
                any_flagged = True
                print(f"FLAGGED: {f['tool']} | Session: {f['sessionId']} | Flags: {', '.join(f['flags'])}")
            else:
                print(f"CLEAN: {f['tool']} | Session: {f['sessionId']}")
        
        print(f"Summary: {len(findings)} calls analyzed, {sum(1 for f in findings if f['flags'])} flagged.")
        
        sys.exit(1 if any_flagged else 0)
        
    except Exception as e: # noqa: BLE001 CLI error handling: broad catch on HTTP requests
        print(f"Error auditing session: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()