from __future__ import annotations

import base64
import json
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from device_auth import DeviceAuthFlow, DeviceAuthSession
from device_auth_protocol import DeviceAuthProtocol, DeviceAuthPromptParts


def _make_fake_jwt(payload: str) -> str:
    encoded = base64.urlsafe_b64encode(payload.encode("utf-8")).decode("ascii").rstrip("=")
    return f"header.{encoded}.signature"


def _build_fake_process(
    stdout_chunks,
    stderr_chunks=(),
    wait_result=0,
    wait_raises=None,
    terminate_changes_returncode=None,
):
    class FakeStream:
        def __init__(self, chunks):
            self._chunks = list(chunks)

        def read(self, _size=-1):
            if not self._chunks:
                return ""
            return self._chunks.pop(0)

    class FakeProcess:
        def __init__(self):
            self.stdout = FakeStream(stdout_chunks)
            self.stderr = FakeStream(stderr_chunks)
            self.returncode = None
            self.terminated = False
            self.wait_calls = []

        def wait(self, timeout=None):
            self.wait_calls.append(timeout)
            if wait_raises is not None:
                raise wait_raises
            self.returncode = wait_result
            return wait_result

        def poll(self):
            return self.returncode

        def terminate(self):
            self.terminated = True
            if terminate_changes_returncode is not None:
                self.returncode = terminate_changes_returncode

        def kill(self):
            self.returncode = -9

    return FakeProcess()


def test_read_prompt_strips_ansi_and_extracts_url_and_code(tmp_path, monkeypatch):
    output = [
        "\x1b[32mOpen this URL to continue:\x1b[0m\n",
        "  \x1b[1mhttps://auth.openai.com/codex/device\x1b[0m\n",
        "Then enter code \x1b[33mH6NU-AJGZ0\x1b[0m to authorize.\n",
    ]
    fake_process = _build_fake_process(output)

    def fake_popen(*args, **kwargs):
        return fake_process

    monkeypatch.setattr(subprocess, "Popen", fake_popen)

    session = DeviceAuthFlow().start(tmp_path, backup_existing=False)
    prompt = DeviceAuthFlow().read_prompt(session, timeout_secs=0.1)

    assert prompt.url == "https://auth.openai.com/codex/device"
    assert prompt.code == "H6NU-AJGZ0"
    assert "Open this URL to continue" in prompt.raw_text
    assert "\x1b[" not in prompt.raw_text


def test_device_auth_protocol_feeds_chunks_and_extracts_prompt() -> None:
    protocol = DeviceAuthProtocol()

    buffered, prompt = protocol.feed("Open ", "\x1b[32mhttps://auth.openai.com/codex/device\x1b[0m\n")
    buffered, prompt = protocol.feed(buffered, "Then enter code H6NU-AJGZ0 to authorize.\n")

    assert buffered == "Open https://auth.openai.com/codex/device\nThen enter code H6NU-AJGZ0 to authorize.\n"
    assert prompt == DeviceAuthPromptParts(
        url="https://auth.openai.com/codex/device",
        code="H6NU-AJGZ0",
        raw_text="Open https://auth.openai.com/codex/device\nThen enter code H6NU-AJGZ0 to authorize.\n",
    )


def test_start_moves_existing_auth_to_backup_before_spawn(tmp_path, monkeypatch):
    auth_path = tmp_path / "auth.json"
    auth_path.write_bytes(b'{"token":"before"}')
    fake_process = _build_fake_process([])
    popen_calls = []

    def fake_popen(*args, **kwargs):
        popen_calls.append((args, kwargs))
        assert not auth_path.exists()
        assert (tmp_path / "auth.json.bak").read_bytes() == b'{"token":"before"}'
        return fake_process

    monkeypatch.setattr(subprocess, "Popen", fake_popen)

    session = DeviceAuthFlow().start(tmp_path, backup_existing=True)

    assert session.had_backup is True
    assert not auth_path.exists()
    assert (tmp_path / "auth.json.bak").read_bytes() == b'{"token":"before"}'
    assert popen_calls


def test_rollback_restores_backup_bytes_exactly(tmp_path):
    original = b'{"id_token":"before","refresh_token":"keep"}'
    (tmp_path / "auth.json.bak").write_bytes(original)
    (tmp_path / "auth.json").write_bytes(b'{"id_token":"after"}')

    DeviceAuthFlow().rollback(tmp_path)

    assert (tmp_path / "auth.json").read_bytes() == original
    assert not (tmp_path / "auth.json.bak").exists()


def test_commit_removes_backup_only(tmp_path):
    auth_bytes = b'{"id_token":"live"}'
    (tmp_path / "auth.json").write_bytes(auth_bytes)
    (tmp_path / "auth.json.bak").write_bytes(b'{"id_token":"stale"}')

    DeviceAuthFlow().commit(tmp_path)

    assert (tmp_path / "auth.json").read_bytes() == auth_bytes
    assert not (tmp_path / "auth.json.bak").exists()


def test_await_completion_times_out_and_terminates_process(tmp_path):
    fake_process = _build_fake_process(
        [], wait_raises=subprocess.TimeoutExpired(cmd="codex", timeout=0.01)
    )

    session = DeviceAuthSession(process=fake_process, codex_home=tmp_path, had_backup=False)
    result = DeviceAuthFlow().await_completion(session, timeout_secs=0.01)

    assert result is False
    assert fake_process.terminated is True


def test_await_completion_succeeds_only_with_decodable_id_token(tmp_path):
    payload = '{"email":"user@example.com"}'
    (tmp_path / "auth.json").write_text(
        '{"id_token":"' + _make_fake_jwt(payload) + '"}', encoding="utf-8"
    )
    fake_process = _build_fake_process([], wait_result=0)

    session = DeviceAuthSession(process=fake_process, codex_home=tmp_path, had_backup=False)
    assert DeviceAuthFlow().await_completion(session, timeout_secs=0.01) is True


def test_await_completion_accepts_current_codex_tokens_shape(tmp_path):
    payload = '{"email":"user@example.com"}'
    auth_payload = {
        "auth_mode": "chatgpt",
        "last_refresh": "2026-07-01T00:00:00Z",
        "tokens": {"id_token": _make_fake_jwt(payload)},
    }
    (tmp_path / "auth.json").write_text(json.dumps(auth_payload), encoding="utf-8")
    fake_process = _build_fake_process([], wait_result=0)

    session = DeviceAuthSession(process=fake_process, codex_home=tmp_path, had_backup=False)
    assert DeviceAuthFlow().await_completion(session, timeout_secs=0.01) is True


def test_start_spawns_resolved_codex_skipping_agent_session_shim(tmp_path, monkeypatch):
    shim_dir = tmp_path / "shim"
    real_dir = tmp_path / "real"
    shim_dir.mkdir()
    real_dir.mkdir()
    shim_body = shim_dir / "_tmpjail-shim.sh"
    shim_body.write_text("#!/bin/sh\n", encoding="utf-8")
    shim_body.chmod(0o755)
    (shim_dir / "codex").symlink_to(shim_body)
    real = real_dir / "codex"
    real.write_text("#!/bin/sh\n", encoding="utf-8")
    real.chmod(0o755)
    monkeypatch.setenv("PATH", f"{shim_dir}:{real_dir}")
    monkeypatch.delenv("CODEX_SQLITE_HOME", raising=False)

    popen_calls = []
    fake_process = _build_fake_process([])

    def fake_popen(*args, **kwargs):
        popen_calls.append((args, kwargs))
        return fake_process

    monkeypatch.setattr(subprocess, "Popen", fake_popen)

    DeviceAuthFlow().start(tmp_path, backup_existing=False)

    argv = popen_calls[0][0][0]
    assert argv[0] == str(real.resolve())
    assert argv[1:] == ["login", "--device-auth"]
    env = popen_calls[0][1]["env"]
    assert env["CODEX_HOME"] == str(tmp_path)
    assert env["CODEX_SQLITE_HOME"] == str(Path.home() / ".codex-shared-state")


def test_start_falls_back_to_codex_when_only_shim_on_path(tmp_path, monkeypatch):
    shim_dir = tmp_path / "shim"
    shim_dir.mkdir()
    shim_body = shim_dir / "_tmpjail-shim.sh"
    shim_body.write_text("#!/bin/sh\n", encoding="utf-8")
    shim_body.chmod(0o755)
    (shim_dir / "codex").symlink_to(shim_body)
    monkeypatch.setenv("PATH", str(shim_dir))

    popen_calls = []
    fake_process = _build_fake_process([])

    def fake_popen(*args, **kwargs):
        popen_calls.append((args, kwargs))
        return fake_process

    monkeypatch.setattr(subprocess, "Popen", fake_popen)

    DeviceAuthFlow().start(tmp_path, backup_existing=False)

    assert popen_calls[0][0][0][0] == "codex"