#!/usr/bin/env python3
"""Headless Codex Integration for Vicoa via ACP."""

import argparse
import logging
import os
import sys
import threading
import uuid
from typing import Any, Dict, Optional

from integrations.headless.acp_base import ACPWrapperBase, ACPWrapperConfig


logger = logging.getLogger(__name__)


class CodexACPConfig(ACPWrapperConfig):
    """Configuration for Codex ACP integration."""

    def __init__(
        self,
        api_key: str,
        base_url: str,
        agent_instance_id: str,
        project_path: str,
        name: str = "Codex",
        is_resuming: bool = False,
        codex_acp_command: str = "codex-acp",
        auth_method: Optional[str] = None,
        codex_api_key: Optional[str] = None,
        openai_api_key: Optional[str] = None,
    ):
        self.api_key = api_key
        self.base_url = base_url
        self.agent_instance_id = agent_instance_id
        self.project_path = project_path
        self.agent_type = "Codex"
        self.agent_command = codex_acp_command
        self.name = name
        self.is_resuming = is_resuming
        self.auth_method = auth_method
        self.codex_api_key = codex_api_key
        self.openai_api_key = openai_api_key

    def get_acp_command(self) -> list[str]:
        return [self.agent_command]

    def get_acp_env(self) -> dict[str, str]:
        env: dict[str, str] = {}
        if self.codex_api_key:
            env["CODEX_API_KEY"] = self.codex_api_key
        if self.openai_api_key:
            env["OPENAI_API_KEY"] = self.openai_api_key
        return env

    @classmethod
    def from_args(
        cls,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        project_path: Optional[str] = None,
        agent_instance_id: Optional[str] = None,
        name: str = "Codex",
        is_resuming: bool = False,
        codex_acp_command: Optional[str] = None,
        auth_method: Optional[str] = None,
        codex_api_key: Optional[str] = None,
        openai_api_key: Optional[str] = None,
    ) -> "CodexACPConfig":
        final_api_key = api_key or os.environ.get("VICOA_API_KEY")
        if not final_api_key:
            raise ValueError(
                "Vicoa API key required: provide --api-key or set VICOA_API_KEY"
            )

        final_base_url = (
            base_url
            or os.environ.get("VICOA_API_URL")
            or os.environ.get("VICOA_BASE_URL")
            or "https://api.vicoa.ai"
        )
        final_project_path = project_path or os.getcwd()
        final_agent_instance_id = (
            agent_instance_id
            or os.environ.get("VICOA_AGENT_INSTANCE_ID")
            or str(uuid.uuid4())
        )

        final_codex_api_key = codex_api_key or os.environ.get("CODEX_API_KEY")
        final_openai_api_key = openai_api_key or os.environ.get("OPENAI_API_KEY")
        final_auth_method = (
            auth_method
            or os.environ.get("VICOA_CODEX_ACP_AUTH_METHOD")
            or (
                "codex-api-key"
                if final_codex_api_key
                else ("openai-api-key" if final_openai_api_key else "chatgpt")
            )
        )

        final_codex_acp_command = (
            codex_acp_command or os.environ.get("VICOA_CODEX_ACP_PATH") or "codex-acp"
        )

        return cls(
            api_key=final_api_key,
            base_url=final_base_url,
            agent_instance_id=final_agent_instance_id,
            project_path=final_project_path,
            name=name,
            is_resuming=is_resuming,
            codex_acp_command=final_codex_acp_command,
            auth_method=final_auth_method,
            codex_api_key=final_codex_api_key,
            openai_api_key=final_openai_api_key,
        )


class CodexACPWrapper(ACPWrapperBase):
    """Headless Codex wrapper using ACP protocol."""

    config: CodexACPConfig
    _permission_wait_poll_interval_seconds: float = 1.0
    _permission_wait_timeout_minutes: int = 15
    _prompt_cancel_grace_period_seconds: float = 10.0
    _permission_cancelled_option_id: str = "__vicoa_cancelled__"

    def __init__(self, config: CodexACPConfig):
        super().__init__(config)
        self._prompt_state_lock = threading.Lock()
        self._prompt_in_flight = False
        self._queued_prompts: list[str] = []
        self._prompt_cancel_event = threading.Event()
        self._interrupt_active = False
        self._permission_request_active = False

    def create_session(self) -> None:
        acp = self.acp
        if not acp:
            raise RuntimeError("ACP client not started")

        # codex-acp requires explicit authenticate before new_session.
        auth_payload = {"methodId": self.config.auth_method}
        auth_response = acp.send_request("authenticate", auth_payload, timeout=120.0)
        auth_response.raise_for_error()

        params: Dict[str, Any] = {
            "cwd": self.config.project_path,
            "mcpServers": [],
        }
        response = acp.send_request("session/new", params)
        response.raise_for_error()
        if not response.result:
            raise RuntimeError("session/new response missing result")

        self.session_id = response.result.get("sessionId")
        if not self.session_id:
            raise RuntimeError("session/new response missing sessionId")

        self.log(f"[ACP] Session created: {self.session_id}")

    def send_prompt(self, message: str) -> None:
        if not self.acp or not self.session_id:
            self.log("[WARNING] Cannot send message: ACP not ready")
            return

        if self._handle_control_command(message):
            return

        with self._prompt_state_lock:
            if self._prompt_in_flight:
                self._queued_prompts.append(message)
                self.log("[ACP] Prompt queued while another prompt is running")
                return
            self._prompt_in_flight = True
            self._prepare_for_new_prompt()

        worker = threading.Thread(
            target=self._run_prompt_request,
            args=(message,),
            daemon=True,
        )
        worker.start()

    def _run_prompt_request(self, message: str) -> None:
        try:
            acp = self.acp
            if not acp or not self.session_id:
                raise RuntimeError("ACP client not ready")

            payload = {
                "sessionId": self.session_id,
                "prompt": [{"type": "text", "text": message}],
            }
            response = acp.send_request(
                "session/prompt",
                payload,
                timeout=3600.0,
                cancel_event=self._prompt_cancel_event,
                cancel_grace_period=self._prompt_cancel_grace_period_seconds,
            )
            response.raise_for_error()
        except Exception as e:
            self.log(f"[ERROR] Failed to send prompt: {e}")
            self._send_feedback_message(f"Prompt failed: {e}")
        finally:
            self._flush_assistant_chunk_buffer()
            self._set_awaiting_input_state()

            next_prompt: Optional[str] = None
            with self._prompt_state_lock:
                self._prompt_in_flight = False
                if self._queued_prompts:
                    next_prompt = self._queued_prompts.pop(0)
                    self._prompt_in_flight = True
                    self._prepare_for_new_prompt()

            if next_prompt:
                worker = threading.Thread(
                    target=self._run_prompt_request,
                    args=(next_prompt,),
                    daemon=True,
                )
                worker.start()

    def _prepare_for_new_prompt(self) -> None:
        self._prompt_cancel_event.clear()
        self._interrupt_active = False

    def handle_notification(self, method: str, params: Dict[str, Any]) -> None:
        if method in {"permission/request", "session/request_permission"}:
            self._handle_permission_request(params)
            return
        self.log(f"[ACP] Unhandled notification: {method}")

    def handle_request(self, method: str, params: Dict[str, Any]) -> Dict[str, Any]:
        if method not in {"permission/request", "session/request_permission"}:
            self.handle_notification(method, params)
            return {}

        option_id = self._handle_permission_request(params)
        if option_id == self._permission_cancelled_option_id:
            return {"outcome": {"outcome": "cancelled"}}

        return {"outcome": {"outcome": "selected", "optionId": option_id}}

    def _handle_permission_request(self, params: Dict[str, Any]) -> str:
        tool_call = params.get("toolCall", {})
        options = params.get("options", [])
        title = (
            str(tool_call.get("title", "")).strip()
            or str(tool_call.get("kind", "tool")).strip()
            or "tool"
        )
        message = f"Permission required: {title}\n"

        option_lines: list[str] = []
        normalized_options: list[dict[str, str]] = []
        for idx, opt in enumerate(options):
            option_id = str(opt.get("optionId", "")).strip()
            option_name = (
                str(opt.get("name", "")).strip()
                or str(opt.get("title", "")).strip()
                or option_id
                or f"Option {idx + 1}"
            )
            option_kind = str(opt.get("kind", "")).strip()
            if option_id:
                normalized_options.append(
                    {"option_id": option_id, "name": option_name, "kind": option_kind}
                )
            option_lines.append(f"{idx + 1}. {option_name}")
        if option_lines:
            message += f"\n[OPTIONS]\n{chr(10).join(option_lines)}\n[/OPTIONS]"

        if self._interrupt_active:
            return self._permission_cancelled_option_id

        self._permission_request_active = True
        try:
            return self._wait_for_permission_decision(message, normalized_options)
        finally:
            self._permission_request_active = False

    def _wait_for_permission_decision(
        self, prompt_message: str, options: list[dict[str, str]]
    ) -> str:
        fallback_option = self._select_default_permission_option(
            [{"optionId": opt["option_id"], "kind": opt.get("kind")} for opt in options]
        )
        if not self.vicoa_client:
            return fallback_option

        self._set_agent_status("AWAITING_INPUT")
        self._suspend_vicoa_polling = True

        try:
            if self._interrupt_active:
                return self._permission_cancelled_option_id

            response = self.vicoa_client.send_message(
                content=prompt_message,
                agent_type=self.config.agent_type,
                agent_instance_id=self.config.agent_instance_id,
                requires_user_input=True,
                timeout_minutes=self._permission_wait_timeout_minutes,
                poll_interval=self._permission_wait_poll_interval_seconds,
            )

            permission_message_id = getattr(response, "message_id", None)
            if permission_message_id:
                self.last_message_id = permission_message_id

            queued_messages = list(getattr(response, "queued_user_messages", []) or [])
            for raw_message in queued_messages:
                message = str(raw_message or "").strip()
                if not message:
                    continue

                control = self._parse_control_command(message)
                if control:
                    setting = control.get("setting")
                    if setting == "interrupt":
                        self._handle_interrupt_control()
                        return self._permission_cancelled_option_id
                    self._apply_control_command(control)
                    continue

                selected = self._parse_permission_reply(message, options)
                if selected:
                    return selected

            self._send_feedback_message(
                f"No valid permission response received. Defaulting to '{fallback_option}'."
            )
            return fallback_option
        except Exception as e:
            self.log(
                f"[WARNING] Permission request failed, defaulting to {fallback_option}: {e}"
            )
            return fallback_option
        finally:
            self._suspend_vicoa_polling = False

    def _parse_permission_reply(
        self, message: str, options: list[dict[str, str]]
    ) -> str | None:
        normalized = message.strip().lower()
        if not normalized:
            return None

        if normalized.isdigit():
            idx = int(normalized) - 1
            if 0 <= idx < len(options):
                return options[idx]["option_id"]

        for opt in options:
            option_id = opt["option_id"]
            option_name = opt["name"]
            option_kind = opt.get("kind", "")
            if normalized in {
                option_id.lower(),
                option_name.lower(),
                option_kind.lower(),
            }:
                return option_id
        return None

    def _select_default_permission_option(self, options: list[Dict[str, Any]]) -> str:
        if not options:
            return "abort"

        for opt in options:
            option_id = str(opt.get("optionId") or "").strip()
            option_kind = str(opt.get("kind") or "").strip().lower()
            if option_kind in {"reject_once", "rejectonce"} and option_id:
                return option_id

        first = str(options[0].get("optionId") or "").strip()
        return first or "abort"

    def _handle_interrupt_control(self) -> None:
        self._interrupt_active = True
        self._prompt_cancel_event.set()
        super()._handle_interrupt_control()


def main() -> int:
    parser = argparse.ArgumentParser(description="Headless Codex Integration via ACP")
    parser.add_argument("--api-key", help="Vicoa API key")
    parser.add_argument("--base-url", help="Vicoa base URL")
    parser.add_argument("--project-path", help="Project directory")
    parser.add_argument("--name", default="Codex", help="Agent display name")
    parser.add_argument("--agent-instance-id", help="Agent instance ID")
    parser.add_argument("--resume", help="Resume session ID")
    parser.add_argument(
        "--codex-acp-command",
        default=None,
        help="Path to codex-acp binary (defaults to VICOA_CODEX_ACP_PATH or codex-acp)",
    )
    parser.add_argument(
        "--auth-method",
        choices=["chatgpt", "codex-api-key", "openai-api-key"],
        default=None,
        help="codex-acp auth method",
    )
    parser.add_argument("--codex-api-key", help="CODEX_API_KEY override")
    parser.add_argument("--openai-api-key", help="OPENAI_API_KEY override")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    agent_instance_id = args.resume or args.agent_instance_id
    is_resuming = bool(args.resume)

    try:
        config = CodexACPConfig.from_args(
            api_key=args.api_key,
            base_url=args.base_url,
            project_path=args.project_path,
            agent_instance_id=agent_instance_id,
            name=args.name,
            is_resuming=is_resuming,
            codex_acp_command=args.codex_acp_command,
            auth_method=args.auth_method,
            codex_api_key=args.codex_api_key,
            openai_api_key=args.openai_api_key,
        )
        wrapper = CodexACPWrapper(config)
        return wrapper.run()
    except Exception as e:
        logger.error(f"Failed to start Codex ACP wrapper: {e}")
        import traceback

        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())
