import base64
import json
import os
import selectors
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path

from device_auth_protocol import DeviceAuthProtocol
from health_client import codex_executable


@dataclass
class DeviceAuthPrompt:
    url: str
    code: str
    raw_text: str


@dataclass
class DeviceAuthSession:
    process: subprocess.Popen
    codex_home: Path
    had_backup: bool


class DeviceAuthFlow:
    def __init__(self) -> None:
        self._protocol = DeviceAuthProtocol()

    def start(self, codex_home: Path, backup_existing: bool) -> DeviceAuthSession:
        auth_path = codex_home / "auth.json"
        backup_path = codex_home / "auth.json.bak"
        had_backup = False

        codex_home.mkdir(parents=True, exist_ok=True)
        if backup_existing and auth_path.exists():
            os.replace(auth_path, backup_path)
            had_backup = True

        env = os.environ.copy()
        env["CODEX_HOME"] = str(codex_home)
        env.setdefault(
            "CODEX_SQLITE_HOME",
            os.path.join(str(Path.home()), ".codex-shared-state"),
        )
        executable = codex_executable() or "codex"
        process = subprocess.Popen(
            [executable, "login", "--device-auth"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            stdin=subprocess.DEVNULL,
            env=env,
        )
        for stream in (process.stdout, process.stderr):
            if stream is None:
                continue
            try:
                os.set_blocking(stream.fileno(), False)
            except (AttributeError, OSError):
                continue
        return DeviceAuthSession(process=process, codex_home=codex_home, had_backup=had_backup)

    def read_prompt(self, session: DeviceAuthSession, timeout_secs: float = 20.0) -> DeviceAuthPrompt:
        deadline = time.monotonic() + timeout_secs
        stripped = ""

        while time.monotonic() < deadline:
            chunk = self._read_next_chunk(session.process, deadline)
            if chunk:
                stripped, prompt = self._protocol.feed(stripped, chunk)
                if prompt is not None:
                    return DeviceAuthPrompt(
                        url=prompt.url,
                        code=prompt.code,
                        raw_text=prompt.raw_text,
                    )
                continue

            if session.process.poll() is not None:
                break
            time.sleep(0.01)

        raise TimeoutError("Timed out waiting for device-auth prompt")

    def await_completion(self, session: DeviceAuthSession, timeout_secs: float = 900.0) -> bool:
        try:
            exit_code = session.process.wait(timeout=timeout_secs)
        except subprocess.TimeoutExpired:
            terminate_process(session.process)
            return False

        if exit_code != 0:
            return False
        return self._has_decodable_id_token(session.codex_home / "auth.json")

    def commit(self, codex_home: Path) -> None:
        backup_path = codex_home / "auth.json.bak"
        if backup_path.exists():
            backup_path.unlink()

    def rollback(self, codex_home: Path) -> None:
        auth_path = codex_home / "auth.json"
        backup_path = codex_home / "auth.json.bak"
        if backup_path.exists():
            os.replace(backup_path, auth_path)

    def abort(self, session: DeviceAuthSession, codex_home: Path) -> None:
        terminate_process(session.process)
        self.rollback(codex_home)

    def _read_next_chunk(self, process: subprocess.Popen, deadline: float) -> str:
        streams = []
        for stream_name in ("stdout", "stderr"):
            stream = getattr(process, stream_name, None)
            if stream is not None:
                streams.append(stream)

        selector = selectors.DefaultSelector()
        registered = False
        try:
            for stream in streams:
                try:
                    selector.register(stream, selectors.EVENT_READ)
                    registered = True
                except Exception:
                    continue

            if registered:
                timeout = max(0.0, deadline - time.monotonic())
                events = selector.select(timeout)
                for key, _mask in events:
                    try:
                        data = os.read(key.fileobj.fileno(), 4096)
                    except BlockingIOError:
                        continue
                    if data:
                        return self._decode_chunk(data)
            return self._fallback_read(streams)
        finally:
            selector.close()

    def _fallback_read(self, streams) -> str:
        for stream in streams:
            try:
                data = stream.read(1)
            except Exception:
                continue
            if data:
                return self._decode_chunk(data)
        return ""

    @staticmethod
    def _decode_chunk(data) -> str:
        if isinstance(data, bytes):
            return data.decode("utf-8", errors="replace")
        return data

    def _has_decodable_id_token(self, auth_path: Path) -> bool:
        try:
            payload = json.loads(auth_path.read_text(encoding="utf-8"))
        except (FileNotFoundError, OSError, json.JSONDecodeError):
            return False

        token = payload.get("id_token")
        if not isinstance(token, str):
            tokens = payload.get("tokens")
            if isinstance(tokens, dict):
                token = tokens.get("id_token")
        if not isinstance(token, str):
            return False

        parts = token.split(".")
        if len(parts) != 3:
            return False

        encoded_payload = parts[1]
        padding = "=" * (-len(encoded_payload) % 4)
        try:
            decoded = base64.urlsafe_b64decode(encoded_payload + padding)
            json.loads(decoded.decode("utf-8"))
        except (ValueError, json.JSONDecodeError, UnicodeDecodeError):
            return False
        return True


def terminate_process(
    process: subprocess.Popen,
    wait_timeout_secs: float = 5.0,
    kill_timeout_secs: float = 1.0,
) -> None:
    if process.poll() is not None:
        return

    try:
        process.terminate()
    except OSError:
        return

    try:
        process.wait(timeout=wait_timeout_secs)
        return
    except subprocess.TimeoutExpired:
        pass

    try:
        process.kill()
    except OSError:
        return

    try:
        process.wait(timeout=kill_timeout_secs)
    except subprocess.TimeoutExpired:
        pass