"""
Codex Session Manager - Background process management for Codex agents.

Architecture:
- Each session runs `codex exec --json` in a background subprocess
- Output is streamed to .zwarm/sessions/<session_id>/output.jsonl
- Session metadata stored in meta.json
- Can inject follow-up messages by starting new turns with context
- Config is read from .zwarm/codex.toml and passed via -c overrides
"""

from __future__ import annotations

import json
import os
import signal
import subprocess
import time
import tomllib
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any
from uuid import uuid4


class SessionStatus(str, Enum):
    """Status of a codex session."""

    PENDING = "pending"  # Created but not started
    RUNNING = "running"  # Process is running
    COMPLETED = "completed"  # Process exited successfully
    FAILED = "failed"  # Process exited with error
    KILLED = "killed"  # Manually killed


@dataclass
class SessionMessage:
    """A message in a session's history."""

    role: str  # "user", "assistant", "system", "tool"
    content: str
    timestamp: str = ""
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "role": self.role,
            "content": self.content,
            "timestamp": self.timestamp,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "SessionMessage":
        return cls(
            role=data.get("role", "unknown"),
            content=data.get("content", ""),
            timestamp=data.get("timestamp", ""),
            metadata=data.get("metadata", {}),
        )


@dataclass
class CodexSession:
    """A managed Codex session."""

    id: str
    task: str
    status: SessionStatus
    working_dir: Path
    created_at: str
    updated_at: str
    pid: int | None = None
    exit_code: int | None = None
    model: str = "gpt-5.1-codex-mini"
    turn: int = 1
    messages: list[SessionMessage] = field(default_factory=list)
    token_usage: dict[str, int] = field(default_factory=dict)
    error: str | None = None
    # Source tracking: "user" for direct spawns, "orchestrator:<instance_id>" for delegated
    source: str = "user"
    # Adapter used: "codex", "claude_code", etc.
    adapter: str = "codex"

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "task": self.task,
            "status": self.status.value,
            "working_dir": str(self.working_dir),
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "pid": self.pid,
            "exit_code": self.exit_code,
            "model": self.model,
            "turn": self.turn,
            "messages": [m.to_dict() for m in self.messages],
            "token_usage": self.token_usage,
            "error": self.error,
            "source": self.source,
            "adapter": self.adapter,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "CodexSession":
        return cls(
            id=data["id"],
            task=data["task"],
            status=SessionStatus(data["status"]),
            working_dir=Path(data["working_dir"]),
            created_at=data["created_at"],
            updated_at=data["updated_at"],
            pid=data.get("pid"),
            exit_code=data.get("exit_code"),
            model=data.get("model", "gpt-5.1-codex-mini"),
            turn=data.get("turn", 1),
            messages=[SessionMessage.from_dict(m) for m in data.get("messages", [])],
            token_usage=data.get("token_usage", {}),
            error=data.get("error"),
            source=data.get("source", "user"),
            adapter=data.get("adapter", "codex"),
        )

    @property
    def is_running(self) -> bool:
        """Check if the session process is still running."""
        if self.pid is None:
            return False
        try:
            os.kill(self.pid, 0)  # Signal 0 just checks if process exists
            return True
        except OSError:
            return False

    @property
    def short_id(self) -> str:
        """Get first 8 chars of ID for display."""
        return self.id[:8]

    @property
    def runtime(self) -> str:
        """Get human-readable runtime."""
        created = datetime.fromisoformat(self.created_at)
        now = datetime.now()
        delta = now - created

        if delta.total_seconds() < 60:
            return f"{int(delta.total_seconds())}s"
        elif delta.total_seconds() < 3600:
            return f"{int(delta.total_seconds() / 60)}m"
        else:
            return f"{delta.total_seconds() / 3600:.1f}h"

    @property
    def source_display(self) -> str:
        """Get short display string for source."""
        if self.source == "user":
            return "you"
        elif self.source.startswith("orchestrator:"):
            # Extract instance ID and shorten it
            instance_id = self.source.split(":", 1)[1]
            return f"orch:{instance_id[:4]}"
        else:
            return self.source[:8]


class CodexSessionManager:
    """
    Manages background Codex sessions.

    Sessions are stored in:
    .zwarm/sessions/<session_id>/
        meta.json      - Session metadata
        output.jsonl   - Raw JSONL output from codex exec
        turns/
            turn_1.jsonl
            turn_2.jsonl
            ...
    """

    def __init__(self, state_dir: Path | str = ".zwarm"):
        self.state_dir = Path(state_dir)
        self.sessions_dir = self.state_dir / "sessions"
        self.sessions_dir.mkdir(parents=True, exist_ok=True)

    def _load_codex_config(self) -> dict[str, Any]:
        """
        Load codex.toml from state_dir.

        Returns parsed TOML as dict, or empty dict if not found.
        """
        codex_toml = self.state_dir / "codex.toml"
        if not codex_toml.exists():
            return {}
        try:
            with open(codex_toml, "rb") as f:
                return tomllib.load(f)
        except Exception:
            return {}

    def _build_codex_overrides(self, config: dict[str, Any]) -> list[str]:
        """
        Convert codex.toml config to -c override flags.

        Handles nested sections like [features] and [sandbox_workspace_write].

        Returns list of ["-c", "key=value", "-c", "key2=value2", ...]
        """
        overrides = []

        def add_override(key: str, value: Any):
            """Add a -c override for a key=value pair."""
            if isinstance(value, bool):
                value = "true" if value else "false"
            overrides.extend(["-c", f"{key}={value}"])

        for key, value in config.items():
            if isinstance(value, dict):
                # Nested section like [features] or [sandbox_workspace_write]
                for subkey, subvalue in value.items():
                    add_override(f"{key}.{subkey}", subvalue)
            else:
                # Top-level key
                add_override(key, value)

        return overrides

    def _session_dir(self, session_id: str) -> Path:
        """Get the directory for a session."""
        return self.sessions_dir / session_id

    def _meta_path(self, session_id: str) -> Path:
        """Get the metadata file path for a session."""
        return self._session_dir(session_id) / "meta.json"

    def _output_path(self, session_id: str, turn: int = 1) -> Path:
        """Get the output file path for a session turn."""
        session_dir = self._session_dir(session_id)
        turns_dir = session_dir / "turns"
        turns_dir.mkdir(parents=True, exist_ok=True)
        return turns_dir / f"turn_{turn}.jsonl"

    def _save_session(self, session: CodexSession) -> None:
        """Save session metadata."""
        session.updated_at = datetime.now().isoformat()
        meta_path = self._meta_path(session.id)
        meta_path.parent.mkdir(parents=True, exist_ok=True)
        meta_path.write_text(json.dumps(session.to_dict(), indent=2))

    def _load_session(self, session_id: str) -> CodexSession | None:
        """Load session from disk."""
        meta_path = self._meta_path(session_id)
        if not meta_path.exists():
            return None
        try:
            data = json.loads(meta_path.read_text())
            return CodexSession.from_dict(data)
        except (json.JSONDecodeError, KeyError) as e:
            print(f"Error loading session {session_id}: {e}")
            return None

    def list_sessions(self, status: SessionStatus | None = None) -> list[CodexSession]:
        """List all sessions, optionally filtered by status."""
        sessions = []
        if not self.sessions_dir.exists():
            return sessions

        for session_dir in self.sessions_dir.iterdir():
            if not session_dir.is_dir():
                continue
            session = self._load_session(session_dir.name)
            if session:
                # Update status if process died OR output indicates completion
                # (output check is more reliable than PID check due to PID reuse)
                if session.status == SessionStatus.RUNNING:
                    if (
                        self._is_output_complete(session.id, session.turn)
                        or not session.is_running
                    ):
                        self._update_session_status(session)

                if status is None or session.status == status:
                    sessions.append(session)

        # Sort by created_at descending (newest first)
        sessions.sort(key=lambda s: s.created_at, reverse=True)
        return sessions

    def get_session(self, session_id: str) -> CodexSession | None:
        """Get a session by ID (supports partial ID matching)."""
        # Try exact match first
        session = self._load_session(session_id)
        if session:
            if session.status == SessionStatus.RUNNING:
                if (
                    self._is_output_complete(session.id, session.turn)
                    or not session.is_running
                ):
                    self._update_session_status(session)
            return session

        # Try partial match
        for session_dir in self.sessions_dir.iterdir():
            if session_dir.name.startswith(session_id):
                session = self._load_session(session_dir.name)
                if session:
                    if session.status == SessionStatus.RUNNING:
                        if (
                            self._is_output_complete(session.id, session.turn)
                            or not session.is_running
                        ):
                            self._update_session_status(session)
                    return session

        return None

    def _is_output_complete(self, session_id: str, turn: int) -> bool:
        """
        Check if output file indicates the task completed.

        Looks for completion markers like 'turn.completed' or 'task.completed'
        in the JSONL output. This is more reliable than PID checking.
        """
        output_path = self._output_path(session_id, turn)
        if not output_path.exists():
            return False

        try:
            content = output_path.read_text()
            for line in content.strip().split("\n"):
                if not line.strip():
                    continue
                try:
                    event = json.loads(line)
                    event_type = event.get("type", "")
                    # Check for any completion marker
                    if event_type in (
                        "turn.completed",
                        "task.completed",
                        "completed",
                        "done",
                    ):
                        return True
                    # Also check for error as a form of completion
                    if event_type == "error":
                        return True
                except json.JSONDecodeError:
                    continue
        except Exception:
            pass

        return False

    def _update_session_status(self, session: CodexSession) -> None:
        """Update session status after process completion."""
        # Parse output to determine status
        output_path = self._output_path(session.id, session.turn)
        if output_path.exists():
            messages, usage, error = self._parse_output(output_path)
            session.messages = messages
            session.token_usage = usage

            # Check if we got actual assistant responses
            has_response = any(m.role == "assistant" for m in messages)

            if error and not has_response:
                # Only mark as failed if we have an error AND no response
                session.status = SessionStatus.FAILED
                session.error = error
            elif error and has_response:
                # Got response but also an error (e.g., network disconnect at end)
                # Treat as completed but note the error
                session.status = SessionStatus.COMPLETED
                session.error = f"Completed with error: {error}"
            else:
                session.status = SessionStatus.COMPLETED
        else:
            session.status = SessionStatus.FAILED
            session.error = "No output file found"

        self._save_session(session)

    def start_session(
        self,
        task: str,
        working_dir: Path | None = None,
        model: str | None = None,
        sandbox: str = "workspace-write",
        source: str = "user",
        adapter: str = "codex",
    ) -> CodexSession:
        """
        Start a new Codex session in the background.

        Args:
            task: The task description
            working_dir: Working directory for codex (default: cwd)
            model: Model override (default: from codex.toml or gpt-5.1-codex-mini)
            sandbox: Sandbox mode (ignored if full_danger=true in codex.toml)
            source: Who spawned this session ("user" or "orchestrator:<id>")
            adapter: Which adapter to use ("codex", "claude_code")

        Returns:
            The created session

        Note:
            Settings are read from .zwarm/codex.toml and passed via -c overrides.
            Run `zwarm init` to set up the config.
        """
        session_id = str(uuid4())
        working_dir = working_dir or Path.cwd()
        now = datetime.now().isoformat()

        # Load codex config from .zwarm/codex.toml
        codex_config = self._load_codex_config()

        # Get model from config or use default
        effective_model = model or codex_config.get("model", "gpt-5.1-codex-mini")

        # Check if full_danger mode is enabled
        full_danger = codex_config.get("full_danger", False)

        session = CodexSession(
            id=session_id,
            task=task,
            status=SessionStatus.PENDING,
            working_dir=working_dir,
            created_at=now,
            updated_at=now,
            model=effective_model,
            turn=1,
            messages=[SessionMessage(role="user", content=task, timestamp=now)],
            source=source,
            adapter=adapter,
        )

        # Create session directory
        session_dir = self._session_dir(session_id)
        session_dir.mkdir(parents=True, exist_ok=True)

        # Build command with -c overrides from codex.toml
        # This ensures each .zwarm dir has its own codex config
        cmd = ["codex"]

        # Add -c overrides from codex.toml (excluding special keys we handle separately)
        config_for_overrides = {k: v for k, v in codex_config.items() if k not in ("model", "full_danger")}
        cmd.extend(self._build_codex_overrides(config_for_overrides))

        # Add exec command and flags
        cmd.extend([
            "exec",
            "--json",
            "--skip-git-repo-check",
            "--model",
            effective_model,
            "-C",
            str(working_dir.absolute()),
        ])

        # Full danger mode bypasses all safety controls
        if full_danger:
            cmd.append("--dangerously-bypass-approvals-and-sandbox")

        cmd.extend(["--", task])

        # Start process with output redirected to file
        output_path = self._output_path(session_id, 1)
        output_file = open(output_path, "w")

        proc = subprocess.Popen(
            cmd,
            cwd=working_dir,
            stdout=output_file,
            stderr=subprocess.STDOUT,
            start_new_session=True,  # Detach from parent process group
        )

        session.pid = proc.pid
        session.status = SessionStatus.RUNNING
        self._save_session(session)

        return session

    def inject_message(
        self,
        session_id: str,
        message: str,
    ) -> CodexSession | None:
        """
        Inject a follow-up message into a completed session.

        This starts a new turn with the conversation context.

        Args:
            session_id: Session to continue
            message: The follow-up message

        Returns:
            Updated session or None if not found/not ready
        """
        session = self.get_session(session_id)
        if not session:
            return None

        if session.status == SessionStatus.RUNNING:
            # Can't inject while running - would need to implement
            # a more complex IPC mechanism for that
            return None

        # Build context from previous messages
        context_parts = []
        for msg in session.messages:
            if msg.role == "user":
                context_parts.append(f"USER: {msg.content}")
            elif msg.role == "assistant":
                context_parts.append(f"ASSISTANT: {msg.content}")

        # Create augmented prompt with context
        augmented_task = f"""Continue the following conversation:

{chr(10).join(context_parts)}

USER: {message}

Continue from where you left off, addressing the user's new message."""

        # Start new turn
        session.turn += 1
        now = datetime.now().isoformat()
        session.messages.append(
            SessionMessage(role="user", content=message, timestamp=now)
        )

        # Build command
        # Note: --search is a global flag that must come before 'exec'
        cmd = ["codex"]
        if session.web_search:
            cmd.append("--search")
        cmd.extend(
            [
                "exec",
                "--json",
                "--full-auto",
                "--skip-git-repo-check",
                "--model",
                session.model,
                "-C",
                str(session.working_dir.absolute()),
                "--",
                augmented_task,
            ]
        )

        # Start process
        output_path = self._output_path(session.id, session.turn)
        output_file = open(output_path, "w")

        proc = subprocess.Popen(
            cmd,
            cwd=session.working_dir,
            stdout=output_file,
            stderr=subprocess.STDOUT,
            start_new_session=True,
        )

        session.pid = proc.pid
        session.status = SessionStatus.RUNNING
        self._save_session(session)

        return session

    def kill_session(self, session_id: str, delete: bool = False) -> bool:
        """
        Kill a running session.

        Args:
            session_id: Session to kill
            delete: If True, also delete session data entirely

        Returns True if killed, False if not found or not running.
        """
        session = self.get_session(session_id)
        if not session:
            return False

        if session.pid and session.is_running:
            try:
                # Kill the entire process group
                os.killpg(os.getpgid(session.pid), signal.SIGTERM)
                time.sleep(0.5)

                # Force kill if still running
                if session.is_running:
                    os.killpg(os.getpgid(session.pid), signal.SIGKILL)
            except (OSError, ProcessLookupError):
                pass

        if delete:
            return self.delete_session(session.id)

        session.status = SessionStatus.KILLED
        session.error = "Manually killed"
        self._save_session(session)
        return True

    def delete_session(self, session_id: str) -> bool:
        """
        Delete a session entirely (removes from disk).

        Kills the process first if still running.

        Returns True if deleted, False if not found.
        """
        import shutil

        session = self.get_session(session_id)
        if not session:
            return False

        # Kill if running
        if session.pid and session.is_running:
            try:
                os.killpg(os.getpgid(session.pid), signal.SIGTERM)
                time.sleep(0.3)
                if session.is_running:
                    os.killpg(os.getpgid(session.pid), signal.SIGKILL)
            except (OSError, ProcessLookupError):
                pass

        # Remove session directory
        session_dir = self._session_dir(session.id)
        if session_dir.exists():
            shutil.rmtree(session_dir)
            return True

        return False

    def get_output(self, session_id: str, turn: int | None = None) -> str:
        """Get raw JSONL output for a session."""
        session = self.get_session(session_id)
        if not session:
            return ""

        if turn is None:
            turn = session.turn

        output_path = self._output_path(session.id, turn)
        if not output_path.exists():
            return ""

        return output_path.read_text()

    def get_messages(self, session_id: str) -> list[SessionMessage]:
        """Get parsed messages for a session across all turns."""
        session = self.get_session(session_id)
        if not session:
            return []

        all_messages = []

        # Get messages from each turn
        for turn in range(1, session.turn + 1):
            output_path = self._output_path(session.id, turn)
            if output_path.exists():
                messages, _, _ = self._parse_output(output_path)
                all_messages.extend(messages)

        return all_messages

    def _parse_output(
        self, output_path: Path
    ) -> tuple[list[SessionMessage], dict[str, int], str | None]:
        """
        Parse JSONL output from codex exec.

        Returns:
            (messages, token_usage, error)
        """
        messages: list[SessionMessage] = []
        usage: dict[str, int] = {}
        error: str | None = None

        if not output_path.exists():
            return messages, usage, "Output file not found"

        content = output_path.read_text()

        for line in content.strip().split("\n"):
            if not line.strip():
                continue

            try:
                event = json.loads(line)
            except json.JSONDecodeError:
                continue

            event_type = event.get("type", "")

            if event_type == "item.completed":
                item = event.get("item", {})
                item_type = item.get("type", "")

                if item_type == "agent_message":
                    text = item.get("text", "")
                    if text:
                        messages.append(
                            SessionMessage(
                                role="assistant",
                                content=text,
                                timestamp=datetime.now().isoformat(),
                            )
                        )

                elif item_type == "reasoning":
                    # Could optionally capture reasoning
                    pass

                elif item_type == "function_call":
                    # Track tool calls
                    func_name = item.get("name", "unknown")
                    messages.append(
                        SessionMessage(
                            role="tool",
                            content=f"[Calling: {func_name}]",
                            metadata={"function": func_name},
                        )
                    )

                elif item_type == "function_call_output":
                    output = item.get("output", "")
                    if output and len(output) < 500:
                        messages.append(
                            SessionMessage(
                                role="tool",
                                content=f"[Output]: {output[:500]}",
                            )
                        )

            elif event_type == "turn.completed":
                turn_usage = event.get("usage", {})
                for key, value in turn_usage.items():
                    usage[key] = usage.get(key, 0) + value
                # Compute total_tokens if not present
                if "total_tokens" not in usage:
                    usage["total_tokens"] = usage.get("input_tokens", 0) + usage.get(
                        "output_tokens", 0
                    )

            elif event_type == "error":
                error = event.get("message", str(event))

        return messages, usage, error

    def get_trajectory(
        self, session_id: str, full: bool = False, max_output_len: int = 200
    ) -> list[dict]:
        """
        Get the full trajectory of a session - all steps in order.

        Args:
            session_id: Session to get trajectory for
            full: If True, include full untruncated content
            max_output_len: Max length for outputs when full=False

        Returns a list of step dicts with type, summary, and details.
        This shows the "broad strokes" of what the agent did.
        """
        if full:
            max_output_len = 999999  # Effectively unlimited
        session = self.get_session(session_id)
        if not session:
            return []

        trajectory = []

        for turn in range(1, session.turn + 1):
            output_path = self._output_path(session.id, turn)
            if not output_path.exists():
                continue

            content = output_path.read_text()
            step_num = 0

            for line in content.strip().split("\n"):
                if not line.strip():
                    continue

                try:
                    event = json.loads(line)
                except json.JSONDecodeError:
                    continue

                event_type = event.get("type", "")

                if event_type == "item.completed":
                    item = event.get("item", {})
                    item_type = item.get("type", "")
                    step_num += 1

                    if item_type == "reasoning":
                        text = item.get("text", "")
                        summary_len = max_output_len if full else 100
                        trajectory.append(
                            {
                                "turn": turn,
                                "step": step_num,
                                "type": "reasoning",
                                "summary": text[:summary_len]
                                + ("..." if len(text) > summary_len else ""),
                                "full_text": text if full else None,
                            }
                        )

                    elif item_type == "command_execution":
                        cmd = item.get("command", "")
                        output = item.get("aggregated_output", "")
                        exit_code = item.get("exit_code")
                        # Truncate output
                        output_preview = output[:max_output_len]
                        if len(output) > max_output_len:
                            output_preview += "..."
                        trajectory.append(
                            {
                                "turn": turn,
                                "step": step_num,
                                "type": "command",
                                "command": cmd,
                                "output": output_preview.strip(),
                                "exit_code": exit_code,
                            }
                        )

                    elif item_type == "function_call":
                        func_name = item.get("name", "unknown")
                        args = item.get("arguments", {})
                        args_str = str(args)
                        args_len = max_output_len if full else 100
                        trajectory.append(
                            {
                                "turn": turn,
                                "step": step_num,
                                "type": "tool_call",
                                "tool": func_name,
                                "args_preview": args_str[:args_len]
                                + ("..." if len(args_str) > args_len else ""),
                                "full_args": args if full else None,
                            }
                        )

                    elif item_type == "function_call_output":
                        output = item.get("output", "")
                        output_preview = output[:max_output_len]
                        if len(output) > max_output_len:
                            output_preview += "..."
                        trajectory.append(
                            {
                                "turn": turn,
                                "step": step_num,
                                "type": "tool_output",
                                "output": output_preview,
                            }
                        )

                    elif item_type == "agent_message":
                        text = item.get("text", "")
                        summary_len = max_output_len if full else 200
                        trajectory.append(
                            {
                                "turn": turn,
                                "step": step_num,
                                "type": "message",
                                "summary": text[:summary_len]
                                + ("..." if len(text) > summary_len else ""),
                                "full_text": text if full else None,
                                "full_length": len(text),
                            }
                        )

        return trajectory

    def cleanup_completed(self, keep_days: int = 7) -> int:
        """
        Remove old completed/failed/killed sessions.

        Args:
            keep_days: Keep sessions newer than this many days

        Returns:
            Number of sessions cleaned up
        """
        import shutil
        from datetime import timedelta

        cutoff = datetime.now() - timedelta(days=keep_days)
        cleaned = 0

        for session in self.list_sessions():
            if session.status in (
                SessionStatus.COMPLETED,
                SessionStatus.FAILED,
                SessionStatus.KILLED,
            ):
                created = datetime.fromisoformat(session.created_at)
                if created < cutoff:
                    session_dir = self._session_dir(session.id)
                    if session_dir.exists():
                        shutil.rmtree(session_dir)
                        cleaned += 1

        return cleaned
