
"""
Finite State Machine for Gemini Sub-Agent Execution.
Defines explicit states and transitions for robust task lifecycle management.
"""

import time
import json
import asyncio
import logging
import shutil
import subprocess
from enum import Enum, auto
from pathlib import Path
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, field

from .logger_setup import setup_gemini_logging
from .config import MEMORY_DIR

logger = setup_gemini_logging("geminis-fsm")

class AgentState(Enum):
    PENDING = auto()
    PREPARING_CONTEXT = auto()
    EXECUTING_PROCESS = auto()
    VALIDATING_OUTPUT = auto()
    CLEANUP = auto()
    COMPLETED = auto()
    FAILED = auto()
    VALIDATION_FAILED = auto()
    DETACHED = auto()

@dataclass
class ExecutionResult:
    status: str
    output: Dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None
    exit_code: int = 0
    duration: float = 0.0

class SubAgentFSM:
    """
    Manages the lifecycle of a Gemini sub-agent task using a State Machine pattern.
    """
    
    def __init__(
        self,
        prompt: str,
        context_files: List[str],
        session_dir: Path,
        cmd: List[str],
        timeout: int = 600,
        session_id: str = None,
        tmux_session: str = None
    ):
        self.state = AgentState.PENDING
        self.prompt = prompt
        self.context_files = context_files
        self.session_dir = session_dir
        self.cmd = cmd
        self.timeout = timeout
        self.start_time = 0.0
        self.result = ExecutionResult(status="pending")
        self.session_id = session_id or session_dir.name
        self.tmux_session = tmux_session
        self.registry_path = MEMORY_DIR / "sub_agent_registry.json"

    async def run(self) -> ExecutionResult:
        self.start_time = time.time()
        try:
            await self._transition(AgentState.PREPARING_CONTEXT)
            if not self._validate_context():
                return self._fail("Context validation failed")

            await self._transition(AgentState.EXECUTING_PROCESS)
            stdout, stderr, code = await self._execute_process()
            
            await self._transition(AgentState.VALIDATING_OUTPUT)
            parsed_output = self._parse_output(stdout, stderr, code)
            
            if parsed_output.get("status") == "failed" or code != 0:
                self.result = ExecutionResult(
                    status="failed",
                    output=parsed_output,
                    error=stderr if stderr else parsed_output.get("error"),
                    exit_code=code
                )
                await self._transition(AgentState.FAILED)
            else:
                self.result = ExecutionResult(
                    status="success",
                    output=parsed_output,
                    exit_code=0
                )
                await self._transition(AgentState.COMPLETED)

        except Exception as e:
            logger.error(f"FSM Error: {e}", exc_info=True)
            self.result = ExecutionResult(status="system_error", error=str(e))
            await self._transition(AgentState.FAILED)
            
        finally:
            self.result.duration = time.time() - self.start_time
            await self._transition(AgentState.CLEANUP)
            
        return self.result

    async def _transition(self, new_state: AgentState):
        logger.info(f"State Transition: {self.state.name} -> {new_state.name}")
        self.state = new_state
        self._update_registry()

    def _update_registry(self):
        try:
            registry = {}
            if self.registry_path.exists():
                try:
                    registry = json.loads(self.registry_path.read_text())
                except:
                    pass
            
            registry[self.session_id] = {
                "status": self.state.name,
                "prompt": self.prompt[:100] + "..." if len(self.prompt) > 100 else self.prompt,
                "tmux_session": self.tmux_session,
                "start_time": self.start_time,
                "last_update": time.time(),
                "duration": time.time() - self.start_time if self.start_time > 0 else 0
            }
            
            now = time.time()
            registry = {k: v for k, v in registry.items() if now - v.get("last_update", 0) < 86400}
            
            self.registry_path.write_text(json.dumps(registry, indent=2))
        except Exception as e:
            logger.error(f"Failed to update registry: {e}")

    async def _fail(self, reason: str) -> ExecutionResult:
        self.result = ExecutionResult(status="failed", error=reason)
        await self._transition(AgentState.FAILED)
        return self.result

    def _validate_context(self) -> bool:
        valid = True
        for f in self.context_files:
            if not Path(f).exists():
                logger.warning(f"Missing context file: {f}")
        return valid

    async def _execute_process(self):
        process = await asyncio.create_subprocess_exec(
            *self.cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            cwd=str(self.session_dir)
        )
        
        try:
            stdout_bytes, stderr_bytes = await asyncio.wait_for(process.communicate(), timeout=self.timeout)
            stdout = stdout_bytes.decode(errors='replace').strip()
            stderr = stderr_bytes.decode(errors='replace').strip()
            
            (self.session_dir / "stdout.log").write_text(stdout)
            (self.session_dir / "stderr.log").write_text(stderr)
            
            return stdout, stderr, process.returncode
        except asyncio.TimeoutError:
            process.kill()
            raise TimeoutError(f"Task timed out after {self.timeout}s")

    def _parse_output(self, raw_stdout: str, raw_stderr: str, exit_code: int) -> Dict[str, Any]:
        if not raw_stdout:
            return {
                "status": "failed",
                "error": "Empty output from Gemini CLI",
                "stderr": raw_stderr,
                "exit_code": exit_code
            }

        try:
            json_str = raw_stdout
            if "{" in raw_stdout:
                json_start = raw_stdout.find("{")
                json_end = raw_stdout.rfind("}") + 1
                if json_end > json_start:
                    json_str = raw_stdout[json_start:json_end]
            
            return json.loads(json_str)
        except json.JSONDecodeError:
            try:
                if "```json" in raw_stdout:
                    clean = raw_stdout.split("```json")[1].split("```")[0].strip()
                    return json.loads(clean)
            except (IndexError, json.JSONDecodeError):
                pass

            return {
                "error": "Failed to parse JSON output",
                "raw_stdout": raw_stdout,
                "raw_stderr": raw_stderr,
                "status": "fatal"
            }

    async def cleanup(self):
        if self.session_dir.exists():
            shutil.rmtree(self.session_dir, ignore_errors=True)
