"""Official version-pinned App Server, ephemeral credentials and offline per-model tool gates.

Only named account/plan operations are exposed by HTTP. The internal RPC transport
is not a proxy. A read-only sandbox alone is never accepted as a no-tools certificate.
"""

import asyncio
import json
import os
import re
import secrets
import shutil
import tempfile
from pathlib import Path
from urllib.parse import urlsplit

from ..protocol import DecisionError, loads, output_schema, validate_plan
from .openai import INSTRUCTIONS, context

VERSION = "codex-cli 0.147.0"
ALLOWED = {
    "initialize",
    "account/login/start",
    "account/login/cancel",
    "account/logout",
    "account/read",
    "model/list",
    "account/rateLimits/read",
    "config/read",
    "thread/start",
    "turn/start",
    "turn/interrupt",
    "thread/unsubscribe",
}
CONFIG = """project_doc_max_bytes = 0
web_search = "disabled"
approval_policy = "never"
sandbox_mode = "read-only"
cli_auth_credentials_store = "ephemeral"
[analytics]
enabled = false
[feedback]
enabled = false
[history]
persistence = "none"
[tools.update_plan]
enabled = false
[tools.experimental_request_user_input]
enabled = false
[features]
apps = false
connectors = false
enable_mcp_apps = false
codex_hooks = false
plugin_hooks = false
hooks = false
skill_search = false
code_mode = false
code_mode_only = false
code_mode_host = false
image_generation = false
computer_use = false
browser_use = false
multi_agent_v2 = false
view_image = false
shell_tool = false
unified_exec = false
multi_agent = false
plugins = false
remote_plugin = false
shell_snapshot = false
skill_mcp_dependency_install = false
"""


def turn_error(error):
    info = error.get("codexErrorInfo") if isinstance(error, dict) else None
    codes = {
        "usageLimitExceeded": "codex_usage_limit_exceeded",
        "sessionBudgetExceeded": "codex_session_budget_exceeded",
        "unauthorized": "codex_auth_required",
        "badRequest": "codex_model_or_request_rejected",
        "serverOverloaded": "codex_server_overloaded",
    }
    code = codes.get(info, "codex_turn_failed") if isinstance(info, str) else "codex_turn_failed"
    return DecisionError(code, 502)


class CodexAccount:
    def __init__(self, directory: Path, *, probe_url=None):
        self.directory = directory
        self.probe_url = probe_url  # Internal offline fake server only, never supplied over HTTP.
        self.probe_token = secrets.token_urlsafe(24) if probe_url else None
        self.session_dir = None
        self.cwd = None
        self.queues = {}
        self.checked = set()
        self.process = None
        self.reader = None
        self.pending = {}
        self.serial = 0
        self.login_id = None
        self.login_complete = False
        self.login_results = {}
        self.lock = asyncio.Lock()
        self.account_lock = asyncio.Lock()

    async def start(self):
        async with self.lock:
            if self.process and self.process.returncode is None:
                return
            binary = shutil.which("codex")
            if not binary:
                raise DecisionError("codex_not_installed", 409)
            self.directory.mkdir(mode=0o700, parents=True, exist_ok=True)
            # Fresh directory per process, never reuse even this application's old auth/config.
            if self.session_dir:
                await self._close()
            self.session_dir = tempfile.TemporaryDirectory(prefix="session-", dir=self.directory)
            home = Path(self.session_dir.name) / "home"
            cwd = Path(self.session_dir.name) / "workspace"
            self.cwd = cwd
            home.mkdir(mode=0o700)
            cwd.mkdir(mode=0o700)
            catalog = home / "models.json"
            catalog.write_bytes(Path(__file__).with_name("codex_models_0_147.json").read_bytes())
            config = "model_catalog_json = " + json.dumps(str(catalog)) + "\n"
            if self.probe_url:
                config += 'model_provider = "offline_probe"\n'
            else:
                config += 'model_provider = "openai"\nforced_login_method = "chatgpt"\n'
            config += CONFIG
            if self.probe_url:
                config += (
                    '\n[model_providers.offline_probe]\nname = "Offline gate"\n'
                    "base_url = " + json.dumps(self.probe_url) + "\n"
                    'wire_api = "responses"\nenv_key = "OFFLINE_PROBE_KEY"\n'
                    "request_max_retries = 0\nstream_max_retries = 0\n"
                )
            (home / "config.toml").write_text(config)
            env = {
                "PATH": os.environ.get("PATH", "/usr/bin:/bin"),
                "HOME": str(home),
                "CODEX_HOME": str(home),
                "XDG_CONFIG_HOME": str(home / "config"),
                "XDG_CACHE_HOME": str(home / "cache"),
                "RUST_LOG": "off",
            }
            if self.probe_url:
                env["OFFLINE_PROBE_KEY"] = self.probe_token
            probe = await asyncio.create_subprocess_exec(
                binary,
                "--version",
                env=env,
                cwd=cwd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.DEVNULL,
            )
            try:
                stdout, _ = await asyncio.wait_for(probe.communicate(), 5)
            except BaseException as exc:
                if probe.returncode is None:
                    probe.kill()
                await probe.wait()
                if isinstance(exc, TimeoutError):
                    raise DecisionError("codex_version_timeout", 504) from None
                raise
            if stdout.decode().strip() != VERSION:
                raise DecisionError("codex_version_unsupported", 409)
            self.process = await asyncio.create_subprocess_exec(
                binary,
                "app-server",
                "--strict-config",
                "--listen",
                "stdio://",
                cwd=cwd,
                env=env,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.DEVNULL,
                limit=1048576,
            )
            self.reader = asyncio.create_task(self._read())
            try:
                await self.rpc(
                    "initialize",
                    {
                        "clientInfo": {"name": "mujoco-decision", "version": "1"},
                        "capabilities": {"experimentalApi": False},
                    },
                )
                self.process.stdin.write(b'{"method":"initialized"}\n')
                await self.process.stdin.drain()
                effective = await self.rpc("config/read", {"cwd": str(cwd), "includeLayers": False})
                config_data = effective.get("config", {})
                if any(config_data.get(key) for key in ("mcp_servers", "hooks", "notify")):
                    raise DecisionError("codex_inherited_execution_config", 409)
            except BaseException:
                await self._close()
                raise

    async def _read(self):
        try:
            while line := await self.process.stdout.readline():
                data = json.loads(line)
                if "id" in data and "method" in data:
                    # No server-initiated tool/approval requests are accepted.
                    raise DecisionError("codex_unexpected_server_request", 502)
                future = self.pending.pop(data.get("id"), None)
                if future and not future.done():
                    if "error" in data:
                        future.set_exception(DecisionError("codex_rpc_failed", 502))
                    else:
                        future.set_result(data.get("result", {}))
                params = data.get("params", {})
                queue = self.queues.get(params.get("threadId"))
                if queue and data.get("method") in ("item/completed", "turn/completed", "error"):
                    queue.put_nowait(data)
                if data.get("method") == "account/login/completed":
                    params = data.get("params", {})
                    ident = params.get("loginId")
                    if isinstance(ident, str):
                        self.login_results[ident] = params.get("success") is True
                        if len(self.login_results) > 8:
                            self.login_results.pop(next(iter(self.login_results)))
                    if ident == self.login_id:
                        self.login_complete = params.get("success") is True
                        self.login_id = None
        except (ValueError, OSError, DecisionError, asyncio.QueueFull):
            pass
        finally:
            for future in self.pending.values():
                if not future.done():
                    future.set_exception(DecisionError("codex_process_exited", 502))
            self.pending.clear()
            for queue in self.queues.values():
                if not queue.full():
                    queue.put_nowait({"method": "error", "params": {}})
            if self.process.returncode is None:
                self.process.terminate()

    async def rpc(self, method, params):
        if method not in ALLOWED:
            raise DecisionError("codex_rpc_forbidden", 403)
        if not self.process or self.process.returncode is not None:
            raise DecisionError("codex_not_running", 409)
        self.serial += 1
        ident = self.serial
        future = asyncio.get_running_loop().create_future()
        self.pending[ident] = future
        try:
            self.process.stdin.write(
                (json.dumps({"id": ident, "method": method, "params": params}) + "\n").encode()
            )
            await self.process.stdin.drain()
            return await asyncio.wait_for(future, 15)
        except TimeoutError:
            raise DecisionError("codex_rpc_timeout", 504) from None
        except (BrokenPipeError, ConnectionError):
            raise DecisionError("codex_process_exited", 502) from None
        finally:
            self.pending.pop(ident, None)

    async def status(self):
        await self.start()
        account = (await self.rpc("account/read", {"refreshToken": False})).get("account")
        return {
            "version": VERSION,
            "experimental": True,
            "storage": "session-only",
            "loggedIn": isinstance(account, dict) and account.get("type") == "chatgpt",
            "planningAvailable": bool(self.checked),
            "checkedModels": sorted(self.checked),
        }

    async def login(self, *, device=False):
        async with self.account_lock:
            try:
                return await self._login(device=device)
            except asyncio.CancelledError:
                # A disconnected browser may lose the RPC reply containing loginId.
                # Close this isolated process rather than leave an unknown login alive.
                await self.close()
                raise

    async def _login(self, *, device=False):
        await self.start()
        if self.login_id:
            raise DecisionError("codex_login_pending", 409)
        value = await self.rpc(
            "account/login/start", {"type": "chatgptDeviceCode" if device else "chatgpt"}
        )
        url = value.get("verificationUrl" if device else "authUrl", "")
        if not isinstance(url, str):
            await self.close()
            raise DecisionError("codex_unexpected_login_url", 502)
        parsed = urlsplit(url)
        if (
            parsed.scheme != "https"
            or parsed.hostname != "auth.openai.com"
            or parsed.username
            or parsed.password
            or parsed.port not in (None, 443)
        ):
            await self.close()
            raise DecisionError("codex_unexpected_login_url", 502)
        self.login_id = value.get("loginId")
        if self.login_id in self.login_results:
            self.login_complete = self.login_results.pop(self.login_id)
            self.login_id = None
        if device:
            code = value.get("userCode")
            if not isinstance(code, str) or not re.fullmatch(r"[A-Za-z0-9-]{4,32}", code):
                await self.close()
                raise DecisionError("codex_invalid_device_code", 502)
            return {"verificationUrl": url, "userCode": code, "storage": "session-only"}
        return {"authUrl": url, "storage": "session-only", "planningAvailable": bool(self.checked)}

    async def cancel_login(self):
        async with self.account_lock:
            return await self._cancel_login()

    async def _cancel_login(self):
        if self.login_id:
            try:
                await self.rpc("account/login/cancel", {"loginId": self.login_id})
            finally:
                self.login_id = None
        return {"cancelled": True}

    async def logout(self):
        async with self.account_lock:
            try:
                if self.process and self.process.returncode is None:
                    await self._cancel_login()
                    await self.rpc("account/logout", {})
            finally:
                await self.close()
        return {"loggedIn": False}

    async def models(self):
        await self.start()
        result = await self.rpc("model/list", {"limit": 100, "includeHidden": False})
        return {
            "models": [
                {"id": m["model"], "name": m["displayName"], "default": m["isDefault"]}
                for m in result.get("data", [])
                if not m.get("hidden")
            ],
            "planningAvailable": bool(self.checked),
        }

    async def limits(self):
        await self.start()
        raw = (await self.rpc("account/rateLimits/read", {})).get("rateLimits", {})
        result = {"source": "codex-app-server", "primary": None, "secondary": None}
        for name in ("primary", "secondary"):
            window = raw.get(name) if isinstance(raw, dict) else None
            if isinstance(window, dict):
                result[name] = {
                    k: v
                    for k, v in window.items()
                    if k in ("usedPercent", "resetsAt", "windowDurationMins")
                    and type(v) is int
                    and 0 <= v <= 10**12
                }
        return result

    async def check_model(self, model):
        models = await self.models()
        if model not in {m["id"] for m in models["models"]}:
            raise DecisionError("codex_model_unavailable", 409)
        if model not in self.checked:
            from .codex_gate import verify_no_tools

            await verify_no_tools(self.directory / "gate", model)
            self.checked.add(model)
        return {"model": model, "toolGatePassed": True, **await self.status()}

    async def plan(self, request, model):
        if not (await self.status())["loggedIn"]:
            raise DecisionError("codex_chatgpt_login_required", 409)
        await self.check_model(model)
        thread = await self.rpc(
            "thread/start",
            {
                "cwd": str(self.cwd),
                "ephemeral": True,
                "approvalPolicy": "never",
                "sandbox": "read-only",
                "model": model,
                "modelProvider": "offline_probe" if self.probe_url else "openai",
                "baseInstructions": INSTRUCTIONS,
            },
        )
        thread_id = thread["thread"]["id"]
        queue = asyncio.Queue(maxsize=32)
        self.queues[thread_id] = queue
        turn_id = None
        complete = False
        try:
            turn = await self.rpc(
                "turn/start",
                {
                    "threadId": thread_id,
                    "input": [{"type": "text", "text": context(request)}],
                    "outputSchema": output_schema("Plan"),
                },
            )
            turn_id = turn["turn"]["id"]
            outputs = []
            async with asyncio.timeout(50):
                while True:
                    event = await queue.get()
                    params = event["params"]
                    if params.get("turnId", turn_id) != turn_id:
                        raise DecisionError("codex_stale_turn", 502)
                    if event["method"] == "error":
                        raise turn_error(params.get("error"))
                    if event["method"] == "item/completed":
                        item = params["item"]
                        if item["type"] == "agentMessage":
                            if item.get("phase") != "commentary":
                                outputs.append(item["text"])
                        elif item["type"] not in ("userMessage", "reasoning"):
                            raise DecisionError("codex_tool_output_forbidden", 502)
                    elif event["method"] == "turn/completed":
                        if params["turn"]["id"] != turn_id:
                            raise DecisionError("codex_stale_turn", 502)
                        if params["turn"]["status"] != "completed":
                            raise turn_error(params["turn"].get("error"))
                        complete = True
                        break
            if len(outputs) != 1 or len(outputs[0]) > 65536:
                raise DecisionError("codex_invalid_output", 502)
            return validate_plan(loads(outputs[0]), request["remaining"]), {}
        finally:
            self.queues.pop(thread_id, None)
            if not complete:
                if turn_id:
                    try:
                        async with asyncio.timeout(3):
                            await self.rpc(
                                "turn/interrupt", {"threadId": thread_id, "turnId": turn_id}
                            )
                    except (DecisionError, asyncio.CancelledError, TimeoutError):
                        await self.close()
                else:
                    await self.close()  # Unknown late turn/start cannot remain alive.
            if self.process and self.process.returncode is None:
                try:
                    async with asyncio.timeout(3):
                        await self.rpc("thread/unsubscribe", {"threadId": thread_id})
                except (DecisionError, TimeoutError):
                    await self.close()

    async def close(self):
        async with self.lock:
            await self._close()

    async def _close(self):
        if self.process:
            if self.process.returncode is None:
                self.process.terminate()
            try:
                await asyncio.wait_for(self.process.wait(), 3)
            except TimeoutError:
                self.process.kill()
                await self.process.wait()
        if self.reader:
            self.reader.cancel()
            await asyncio.gather(self.reader, return_exceptions=True)
        self.process = None
        self.reader = None
        self.login_id = None
        self.login_complete = False
        self.checked.clear()
        self.queues.clear()
        self.login_results.clear()
        if self.session_dir:
            self.session_dir.cleanup()
            self.session_dir = None