
"""
Delegator: Core logic for sub-agent task management, context bundling, and execution.
"""

import os
import json
import uuid
import time
import shutil
import asyncio
import logging
import subprocess
from pathlib import Path
from typing import List, Optional, Dict

from .config import (
    WORKSPACE_ROOT, AGENT_DIR, SKILLS_DIR, SESSIONS_DIR, 
    LOG_DIR, RULES_PATH, ACTIVE_CONTEXT_PATH, WORKFLOWS_DIR, PERSONAS_DIR
)
from .metrics_logger import log_metrics
from .logger_setup import setup_gemini_logging

logger = setup_gemini_logging("geminis-delegator")

def validate_identifier(name: str) -> bool:
    """Ensure name is safe (alphanumeric, -, _) to prevent path traversal."""
    if not name: return False
    return all(c.isalnum() or c in "-_" for c in name) and ".." not in name

def auto_bundle_context(
    context_files: List[str],
    include_rules: bool = True,
    include_workflows: List[str] = None,
    include_personas: List[str] = None,
    include_skills: List[str] = None,
    task_prompt: str = ""
) -> List[str]:
    """Automatically bundle essential context files for sub-agent execution."""
    bundled = list(context_files)
    
    if include_rules and RULES_PATH.exists():
        bundled.append(str(RULES_PATH))
    
    if ACTIVE_CONTEXT_PATH.exists():
        bundled.append(str(ACTIVE_CONTEXT_PATH))
    
    plans_dir = WORKSPACE_ROOT / "docs" / "plans"
    if plans_dir.exists():
        keywords = set(task_prompt.lower().split())
        for plan_file in plans_dir.glob("*.md"):
            plan_name = plan_file.stem.lower()
            if any(word in plan_name for word in keywords if len(word) > 3):
                bundled.append(str(plan_file))

    if include_workflows:
        for wf_name in include_workflows:
            if not validate_identifier(wf_name):
                logger.warning(f"Skipping invalid workflow name: {wf_name}")
                continue
            wf_path = WORKFLOWS_DIR / f"{wf_name}.md"
            if wf_path.exists():
                bundled.append(str(wf_path))
    
    if include_personas:
        for persona_name in include_personas:
            if not validate_identifier(persona_name):
                logger.warning(f"Skipping invalid persona name: {persona_name}")
                continue
            persona_path = PERSONAS_DIR / f"{persona_name}.md"
            if persona_path.exists():
                bundled.append(str(persona_path))
    
    if include_skills:
        for skill_name in include_skills:
            if not validate_identifier(skill_name):
                logger.warning(f"Skipping invalid skill name: {skill_name}")
                continue
            skill_path = SKILLS_DIR / skill_name / "SKILL.md"
            if skill_path.exists():
                bundled.append(str(skill_path))
    
    return bundled

async def run_sub_agent(
    prompt: str,
    context_files: List[str],
    approval_mode: str = "auto_edit",
    lane: str = "default",
    background: bool = False,
    tmux: bool = True,
    repomix_targets: List[str] = None
) -> str:
    """Execute the gemini CLI and manage the session lifecycle using FSM."""
    from .execution_state import SubAgentFSM
    
    session_id = f"session_{int(time.time())}_{uuid.uuid4().hex[:8]}"
    session_dir = SESSIONS_DIR / session_id
    session_dir.mkdir(parents=True, exist_ok=True)

    if repomix_targets:
        for i, target in enumerate(repomix_targets):
            target_path = Path(target)
            if not target_path.is_absolute():
                target_path = WORKSPACE_ROOT / target_path
            
            if target_path.exists():
                output_name = f"context_pack_{i}_{target_path.name}.md"
                output_path = session_dir / output_name
                logger.info(f"Running Repomix on {target_path} -> {output_path}")
                
                try:
                    repomix_args = [
                        "npx", "-y", "repomix", str(target_path),
                        "--style", "markdown",
                        "--output", str(output_path),
                    ]
                    if ".agent/knowledge" in str(target_path):
                        repomix_args.append("--no-gitignore")
                    
                    process = await asyncio.create_subprocess_exec(
                        *repomix_args,
                        stdout=asyncio.subprocess.PIPE,
                        stderr=asyncio.subprocess.PIPE
                    )
                    await process.communicate()
                    
                    if output_path.exists():
                        context_files.append(str(output_path))
                        logger.info(f"Successfully packed context: {output_path}")
                except Exception as e:
                    logger.error(f"Repomix failed for {target}: {e}")

    def sanitize(text: str) -> str:
        return "".join(ch for ch in text if ch.isprintable() or ch in '\n\r\t')
    
    prompt = sanitize(prompt)

    gemini_path = shutil.which("gemini")
    if not gemini_path:
        logger.error("Gemini CLI not found in PATH")
        return json.dumps({
            "status": "failed", 
            "error": "Gemini CLI not found. Please install it or check PATH.",
            "context": {"lane": lane}
        })

    system_intro = (
        "You are a recursive Gemini sub-agent. You have full agency to build, research, and execute.\n"
        "RECURSION ENABLED: You can trigger further sub-tasks by running the mcp tool 'delegate_task'.\n"
        f"ISOLATION: You are running in a temporary session directory: {session_dir}\n"
        f"REPO ROOT: The repository root is at: {WORKSPACE_ROOT}\n"
        "IMPORTANT: The current directory is EMPTY. To read/edit project files, use absolute paths or join with REPO ROOT.\n"
        "\n"
        "*** CRITICAL RULES ***\n"
        "1. ARTIFACTS: ALL plans, reports, todos, and walkthroughs MUST be saved to 'docs/' (e.g., docs/plans/, docs/reports/).\n"
        "2. METRICS: Your usage is automatically logged. Be efficient.\n"
        "Just do it. Don't ask for permission.\n\n"
    )
    
    full_prompt = system_intro + prompt
    valid_context = []
    for file_path in context_files:
        if os.path.exists(file_path):
            full_prompt += f" @{file_path}"
            valid_context.append(file_path)
        else:
            logger.warning(f"Context file not found: {file_path}")

    cmd = [
        "gemini", 
        "--prompt", full_prompt, 
        "--output-format", "json", 
        "--approval-mode", approval_mode,
        "--sandbox", "false",
        "--include-directories", str(WORKSPACE_ROOT)
    ]

    is_tmux_available = shutil.which("tmux") is not None
    tmux_session = None
    if background and tmux and is_tmux_available:
        tmux_session = f"gemini_{session_id[-8:]}"
        log_file = session_dir / "stdout.log"
        monitor_cmd = f"until [ -f {log_file} ]; do sleep 0.1; done; tail -f {log_file}"
        subprocess.Popen(["tmux", "new-session", "-d", "-s", tmux_session, monitor_cmd])
        logger.info(f"Background task monitor started in tmux session: {tmux_session}")
    
    fsm = SubAgentFSM(
        prompt=prompt,
        context_files=valid_context,
        session_dir=session_dir,
        cmd=cmd,
        session_id=session_id,
        tmux_session=tmux_session
    )
    
    if background:
        async def background_wrapper():
            try:
                result = await fsm.run()
                output_data = result.output
                output_data["status"] = result.status
                output_data["context"] = {
                    "files_count": len(valid_context),
                    "lane": lane,
                    "duration_ms": int(result.duration * 1000)
                }
                log_metrics(prompt, output_data)
                if result.status == "success" and session_dir.exists():
                    shutil.rmtree(session_dir, ignore_errors=True)
            except Exception as e:
                logger.error(f"Background task {session_id} failed: {e}")

        asyncio.create_task(background_wrapper())
        
        return json.dumps({
            "status": "detached",
            "session_id": session_id,
            "tmux_session": tmux_session,
            "monitor_cmd": f"tmux attach -t {tmux_session}" if tmux_session else f"tail -f {LOG_DIR}/geminis_mcp.log",
            "context": {
                "files_count": len(valid_context),
                "lane": lane
            }
        })

    result = await fsm.run()
    output_data = result.output
    output_data["status"] = result.status
    output_data["context"] = {
        "files_count": len(valid_context),
        "lane": lane,
        "duration_ms": int(result.duration * 1000)
    }
    log_metrics(prompt, output_data)
    if result.status == "success" and session_dir.exists():
        shutil.rmtree(session_dir, ignore_errors=True)

    return json.dumps(output_data)
