import argparse
import subprocess
import json
import os
import tomllib
import tree_sitter
import tree_sitter_python
import ast
import datetime

LOG_FILE: str = 'logs/errors/mma_delegation.log'


def generate_skeleton(code: str) -> str:
 """
 Parses Python code and replaces function/method bodies with '...',
 preserving docstrings if present.
 """
 try:
  PY_LANGUAGE = tree_sitter.Language(tree_sitter_python.language())
  parser = tree_sitter.Parser(PY_LANGUAGE)
  tree = parser.parse(bytes(code, "utf8"))
  edits = []

  def is_docstring(node):
   if node.type == "expression_statement" and node.child_count > 0:
    if node.children[0].type == "string":
     return True
   return False

  def walk(node):
   if node.type == "function_definition":
    body = node.child_by_field_name("body")
    if body and body.type == "block":
     indent = " " * body.start_point.column
     first_stmt = None
     for child in body.children:
      if child.type != "comment":
       first_stmt = child
       break
     if first_stmt and is_docstring(first_stmt):
      start_byte = first_stmt.end_byte
      end_byte = body.end_byte
      if end_byte > start_byte:
       edits.append((start_byte, end_byte, f"\n{indent}..."))
     else:
      start_byte = body.start_byte
      end_byte = body.end_byte
      edits.append((start_byte, end_byte, "..."))
   for child in node.children:
    walk(child)
  walk(tree.root_node)
  edits.sort(key=lambda x: x[0], reverse=True)
  code_bytes = bytearray(code, "utf8")
  for start, end, replacement in edits:
   code_bytes[start:end] = bytes(replacement, "utf8")
  return code_bytes.decode("utf8")
 except Exception as e:
  return f"# Error generating skeleton: {e}\n{code}"

def get_model_for_role(role: str, failure_count: int = 0) -> str:
 """Returns the specific model to use for a given tier role."""
 if role == 'tier1-orchestrator' or role == 'tier1':
  return 'gemini-3.1-pro-preview'
 elif role == 'tier2-tech-lead' or role == 'tier2':
  return 'gemini-3-flash-preview'
 elif role == 'tier3-worker' or role == 'tier3':
  return 'gemini-3-flash-preview'
 elif role == 'tier4-qa' or role == 'tier4':
  return 'gemini-2.5-flash-lite'
 else:
  return 'gemini-2.5-flash-lite'

def get_role_documents(role: str) -> list[str]:
 if role == 'tier1-orchestrator' or role == 'tier1':
  return ['conductor/product.md', 'conductor/product-guidelines.md',
   'docs/guide_architecture.md', 'docs/guide_mma.md']
 elif role == 'tier2-tech-lead' or role == 'tier2':
  return ['conductor/tech-stack.md', 'conductor/workflow.md',
   'docs/guide_architecture.md', 'docs/guide_mma.md']
 elif role == 'tier3-worker' or role == 'tier3':
  return ['docs/guide_architecture.md']
 elif role == 'tier4-qa' or role == 'tier4':
  return ['docs/guide_architecture.md']
 return []

def log_delegation(role: str, full_prompt: str, result: str | None = None, summary_prompt: str | None = None) -> str:
 os.makedirs('logs/agents', exist_ok=True)
 timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
 log_file = f'logs/agents/mma_{role}_task_{timestamp}.log'
 with open(log_file, 'w', encoding='utf-8') as f:
  f.write("==================================================\n")
  f.write(f"ROLE: {role}\n")
  f.write(f"TIMESTAMP: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
  f.write("--------------------------------------------------\n")
  f.write(f"FULL PROMPT:\n{full_prompt}\n")
  f.write("--------------------------------------------------\n")
  if result:
   f.write(f"RESULT:\n{result}\n")
  f.write("==================================================\n")
  # Also keep the master log
 os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)
 display_prompt = summary_prompt if summary_prompt else full_prompt
 with open(LOG_FILE, 'a', encoding='utf-8') as f:
  f.write(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {role}: {display_prompt[:100]}... (Log: {log_file})\n")
 return log_file

def get_dependencies(filepath: str) -> list[str]:
 """Identify top-level module imports from a Python file."""
 try:
  with open(filepath, 'r', encoding='utf-8') as f:
   tree = ast.parse(f.read())
  dependencies = []
  for node in tree.body:
   if isinstance(node, ast.Import):
    for alias in node.names:
     dependencies.append(alias.name.split('.')[0])
   elif isinstance(node, ast.ImportFrom):
    if node.module:
     dependencies.append(node.module.split('.')[0])
  seen = set()
  result = []
  for d in dependencies:
   if d not in seen:
    result.append(d)
    seen.add(d)
  return result
 except Exception as e:
  print(f"Error getting dependencies for {filepath}: {e}")
  return []

def execute_agent(role: str, prompt: str, docs: list[str], debug: bool = False, failure_count: int = 0, tool_preset: str | None = None) -> str:
 model = get_model_for_role(role, failure_count)

 # Advanced Context: Dependency skeletons for Tier 3
 injected_context = ""
 if role in ['tier3', 'tier3-worker']:
  for doc in docs:
   if doc.endswith('.py') and os.path.exists(doc):
    deps = get_dependencies(doc)
    for dep in deps:
    # Only try to generate skeletons for files that exist in the local dir
     dep_file = f"{dep}.py"
     # Optimization: If the dependency is already in 'docs' (explicitly provided), 
     # do NOT inject its skeleton/full context again as a dependency.
     if dep_file in docs:
      continue
     if os.path.exists(dep_file) and dep_file != doc:
      try:
       with open(dep_file, 'r', encoding='utf-8') as f:
        skeleton = generate_skeleton(f.read())
       injected_context += f"\n\nDEPENDENCY SKELETON: {dep_file}\n{skeleton}\n"
      except Exception as e:
       print(f"Error gathering context for {dep_file}: {e}")
       # Check for token-bloat safety: if injected_context is too large, truncate it
    if len(injected_context) > 15000:
     injected_context = injected_context[:15000] + "... [TRUNCATED FOR COMMAND LINE LIMITS]"
     # MMA Protocol: Tier 3 and 4 are stateless.
 if role in ['tier3', 'tier3-worker']:
  system_directive = "STRICT SYSTEM DIRECTIVE: You are a stateless Tier 3 Worker (Contributor). " \
  "Your goal is to implement specific code changes or tests based on the provided task. " \
  "CRITICAL CODE STYLE RULE: ALL Python code MUST use exactly 1 SPACE for indentation. DO NOT use 4 spaces or tabs. " \
  "You have access to tools for reading and writing files (e.g., read_file, write_file, replace), " \
  "codebase investigation (discovered_tool_py_get_code_outline, discovered_tool_py_get_skeleton, discovered_tool_py_get_definition, discovered_tool_py_find_usages, discovered_tool_py_get_imports, discovered_tool_py_check_syntax, discovered_tool_get_tree), " \
  "version control (discovered_tool_get_git_diff), and web tools (discovered_tool_web_search, discovered_tool_fetch_url). " \
  "You CAN execute PowerShell scripts via discovered_tool_run_powershell for verification and testing. " \
  "ARCHITECTURE REFERENCE: docs/guide_architecture.md contains the threading model, cross-thread data structures, " \
  "frame-sync mechanism (_process_pending_gui_tasks action catalog), AI client architecture, and HITL Execution Clutch. " \
  "CRITICAL THREADING RULE: NEVER write GUI state from a background thread. Push task dicts to _pending_gui_tasks with the lock. " \
  "TASK FORMAT: Your task will specify WHERE (file:line), WHAT (change), HOW (API calls), and SAFETY (thread constraints). Follow these exactly. " \
  "Follow TDD and return success status or code changes. No pleasantries, no conversational filler."
 elif role in ['tier4', 'tier4-qa']:
  system_directive = "STRICT SYSTEM DIRECTIVE: You are a stateless Tier 4 QA Agent. " \
  "Your goal is to analyze errors, summarize logs, or verify tests. " \
  "You have access to tools for reading files, exploring the codebase (discovered_tool_py_get_code_outline, discovered_tool_py_get_skeleton, discovered_tool_py_get_definition, discovered_tool_py_find_usages, discovered_tool_py_get_imports), " \
  "version control (discovered_tool_get_git_diff), and web tools (discovered_tool_web_search, discovered_tool_fetch_url). " \
  "You CAN execute PowerShell scripts via discovered_tool_run_powershell for diagnostics. " \
  "ARCHITECTURE REFERENCE: docs/guide_architecture.md contains the threading model and data flow. " \
  "When analyzing errors, trace the data flow through the thread domains (GUI main, asyncio worker, HookServer) documented there. " \
  "ONLY output the requested analysis. No pleasantries."
 else:
  system_directive = f"STRICT SYSTEM DIRECTIVE: You are a stateless {role}. " \
  "ONLY output the requested text. No pleasantries."
 command_text = f"{system_directive}\n\n{injected_context}\n\n"
 # Manually inline documents to ensure sub-agent has context in headless mode
 for doc in docs:
  if os.path.exists(doc):
   try:
    with open(doc, 'r', encoding='utf-8') as f:
     content = f.read()
    command_text += f"\n\nFILE CONTENT: {doc}\n{content}\n"
   except Exception as e:
    print(f"Error inlining {doc}: {e}")
 command_text += f"\n\nTASK: {prompt}\n\n"
 # Use subprocess with input to pipe the prompt via stdin, avoiding WinError 206.
 # We use -p 'mma_task' to ensure non-interactive (headless) mode and valid parsing.
 ps_command = f"gemini -p '{role}' --output-format json --model {model}"
 cmd = ['powershell.exe', '-NoProfile', '-Command', ps_command]
 try:
  env = os.environ.copy()
  env["GEMINI_CLI_HOOK_CONTEXT"] = "mma_headless"
  if tool_preset is not None:
   env["SLOP_TOOL_PRESET"] = tool_preset
  if debug:
   print("--- MMA DEBUG ---")
   print(f"Executing Command: {cmd}")
   print("Relevant Environment Variables:")
   for key, value in env.items():
    if key.startswith("GEMINI_CLI_") or key == "SLOP_TOOL_PRESET":
     print(f"  {key}={value}")
  process = subprocess.run(cmd, input=command_text, capture_output=True, text=True, encoding='utf-8', env=env)
  if debug:
   print("Subprocess Result:")
   print(f"  Return Code: {process.returncode}")
   print(f"  Stdout: {process.stdout[:1000]}...")
   print(f"  Stderr: {process.stderr}")
   print("--- END DEBUG ---")
  result = process.stdout
  if not process.stdout and process.stderr:
   result = f"Error: {process.stderr}"
   # Log the attempt and result
  log_file = log_delegation(role, command_text, result, summary_prompt=prompt)
  print(f"Sub-agent log created: {log_file}")
  stdout = process.stdout
  start_index = stdout.find('{')
  if start_index != -1:
   json_str = stdout[start_index:]
   try:
    data = json.loads(json_str)
    return data.get('response', stdout)
   except json.JSONDecodeError:
    return stdout
  return stdout
 except Exception as e:
  err_msg = f"Execution failed: {str(e)}"
  log_delegation(role, command_text, err_msg)
  return err_msg

def create_parser() -> argparse.ArgumentParser:
 parser = argparse.ArgumentParser(description="MMA Execution Script")
 parser.add_argument(
  "--role", 
  choices=['tier1', 'tier2', 'tier3', 'tier4', 'tier1-orchestrator', 'tier2-tech-lead', 'tier3-worker', 'tier4-qa'], 
  help="The tier role to execute"
 )
 parser.add_argument(
  "--task-file",
  type=str,
  help="TOML file defining the task"
 )
 parser.add_argument(
  "--debug",
  action="store_true",
  help="Enable debug logging"
 )
 parser.add_argument(
  "--failure-count",
  type=int,
  default=0,
  help="Number of times this task has failed previously"
 )
 parser.add_argument(
  "--tool-preset",
  type=str,
  help="The tool preset to use"
 )
 parser.add_argument(
  "prompt", 
  type=str, 
  nargs='?',
  help="The prompt for the tier (optional if --task-file is used)"
 )
 return parser

def main() -> None:
 parser = create_parser()
 args = parser.parse_args()
 role = args.role
 prompt = args.prompt
 debug = args.debug
 failure_count = args.failure_count
 tool_preset = args.tool_preset
 docs = []
 if args.task_file and os.path.exists(args.task_file):
  with open(args.task_file, "rb") as f:
   task_data = tomllib.load(f)
   role = task_data.get("role", role)
   prompt = task_data.get("prompt", prompt)
   docs = task_data.get("docs", [])
   # Only override debug if it's explicitly set in the task file (optional)
   debug = task_data.get("debug", debug)
   failure_count = task_data.get("failure_count", failure_count)
   tool_preset = task_data.get("tool_preset", tool_preset)
 if not role or not prompt:
  parser.print_help()
  return
 if not docs:
  docs = get_role_documents(role)
  # Extract @file references from the prompt
 import re
 file_refs = re.findall(r"@([\w./\\]+)", prompt)
 for ref in file_refs:
  if os.path.exists(ref) and ref not in docs:
   docs.append(ref)
 print(f"Executing role: {role} with docs: {docs} (debug={debug}, failure_count={failure_count}, tool_preset={tool_preset})")
 result = execute_agent(role, prompt, docs, debug=debug, failure_count=failure_count, tool_preset=tool_preset)
 print(result)

if __name__ == "__main__":
 main()
