from __future__ import annotations

import base64
import json
import sys
from pathlib import Path

import pytest

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

from account_registry import AccountRegistry
from pi_auth_sync import PI_PROVIDER_KEY, PiAuthSync, PiAuthSyncError

EXP_2100 = 4102444800


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


def write_codex_auth(
    path: Path,
    email: str,
    account_id: str,
    *,
    exp: int = EXP_2100,
    include_account_id_field: bool = True,
) -> None:
    claims = {"https://api.openai.com/auth": {"chatgpt_account_id": account_id}}
    tokens = {
        "id_token": _jwt({"email": email, **claims}),
        "access_token": _jwt({"exp": exp, **claims}),
        "refresh_token": f"rt-{account_id}",
    }
    if include_account_id_field:
        tokens["account_id"] = account_id
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps({"tokens": tokens}), encoding="utf-8")


@pytest.fixture
def registry(tmp_path: Path) -> AccountRegistry:
    legacy = tmp_path / "legacy"
    legacy.mkdir()
    pi_agent = tmp_path / "pi-agent"
    pi_agent.mkdir()
    return AccountRegistry(
        base_dir=tmp_path / "tray",
        legacy_codex_home=legacy,
        pi_agent_dir=pi_agent,
    )


def _add_account(registry: AccountRegistry, slug: str, alias: str) -> Path:
    codex_home = registry.add_dir(slug, alias)
    write_codex_auth(codex_home / "auth.json", f"{slug}@example.com", f"acct-{slug}")
    return codex_home


def test_set_default_switches_codex_and_pi_together(registry: AccountRegistry) -> None:
    _add_account(registry, "rafa", "Rafa")
    _add_account(registry, "roy", "Roy")
    pi_link = registry._pi_auth_sync.active_auth_path
    codex_link = registry.legacy_codex_home / "auth.json"

    for slug in ("rafa", "roy", "rafa"):
        account = next(a for a in registry.list() if a.slug == slug)
        registry.set_default(account)

        expected_pi = registry.accounts_dir / slug / "PI_HOME" / "auth.json"
        assert codex_link.resolve() == (
            registry.accounts_dir / slug / "CODEX_HOME" / "auth.json"
        ).resolve()
        assert pi_link.is_symlink()
        assert pi_link.resolve() == expected_pi.resolve()
        credential = json.loads(expected_pi.read_text(encoding="utf-8"))[PI_PROVIDER_KEY]
        assert credential == {
            "type": "oauth",
            "access": json.loads(
                (registry.accounts_dir / slug / "CODEX_HOME" / "auth.json").read_text()
            )["tokens"]["access_token"],
            "refresh": f"rt-acct-{slug}",
            "expires": EXP_2100 * 1000,
            "accountId": f"acct-{slug}",
        }


def test_prepare_merges_into_existing_pi_auth_preserving_other_providers(
    registry: AccountRegistry,
) -> None:
    _add_account(registry, "rafa", "Rafa")
    pi_auth = registry.accounts_dir / "rafa" / "PI_HOME" / "auth.json"
    pi_auth.parent.mkdir(parents=True)
    pi_auth.write_text(json.dumps({"google": {"type": "api-key", "key": "g"}}), encoding="utf-8")

    registry.set_default(registry.list()[0])

    merged = json.loads(pi_auth.read_text(encoding="utf-8"))
    assert merged["google"] == {"type": "api-key", "key": "g"}
    assert merged[PI_PROVIDER_KEY]["accountId"] == "acct-rafa"


def test_missing_pi_agent_dir_skips_pi_with_warning_and_switches_codex(
    tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
    legacy = tmp_path / "legacy"
    legacy.mkdir()
    registry = AccountRegistry(
        base_dir=tmp_path / "tray",
        legacy_codex_home=legacy,
        pi_agent_dir=tmp_path / "absent-pi-agent",
    )
    _add_account(registry, "rafa", "Rafa")

    registry.set_default(registry.list()[0])

    assert (legacy / "auth.json").is_symlink()
    assert not (tmp_path / "absent-pi-agent").exists()
    assert "absent-pi-agent" in capsys.readouterr().err


def test_underivable_codex_auth_fails_closed_without_half_switch(
    registry: AccountRegistry,
) -> None:
    _add_account(registry, "rafa", "Rafa")
    roy_home = registry.add_dir("roy", "Roy")
    (roy_home / "auth.json").write_text(json.dumps({"tokens": {}}), encoding="utf-8")
    registry.set_default(next(a for a in registry.list() if a.slug == "rafa"))
    codex_link = registry.legacy_codex_home / "auth.json"
    pi_link = registry._pi_auth_sync.active_auth_path

    with pytest.raises(PiAuthSyncError):
        registry.set_default(next(a for a in registry.list() if a.slug == "roy"))

    assert registry.default_path.read_text(encoding="utf-8") == "rafa\n"
    assert codex_link.resolve() == (
        registry.accounts_dir / "rafa" / "CODEX_HOME" / "auth.json"
    ).resolve()
    assert pi_link.resolve() == (
        registry.accounts_dir / "rafa" / "PI_HOME" / "auth.json"
    ).resolve()


def test_remove_last_account_clears_pi_link(registry: AccountRegistry) -> None:
    _add_account(registry, "rafa", "Rafa")
    registry.set_default(registry.list()[0])
    pi_link = registry._pi_auth_sync.active_auth_path
    assert pi_link.is_symlink()

    registry.remove("rafa")

    assert not pi_link.is_symlink()
    assert not (registry.legacy_codex_home / "auth.json").is_symlink()


def test_registry_without_pi_agent_dir_has_pi_sync_disabled(tmp_path: Path) -> None:
    legacy = tmp_path / "legacy"
    legacy.mkdir()
    registry = AccountRegistry(base_dir=tmp_path / "tray", legacy_codex_home=legacy)
    assert registry._pi_auth_sync is None


def test_derive_uses_jwt_claim_when_account_id_field_missing(tmp_path: Path) -> None:
    sync = PiAuthSync(agent_dir=tmp_path / "pi", accounts_dir=tmp_path / "accounts")
    auth = tmp_path / "accounts" / "x" / "CODEX_HOME" / "auth.json"
    write_codex_auth(auth, "x@example.com", "acct-x", include_account_id_field=False)

    credential = sync._derive_credential(auth, "x")

    assert credential["accountId"] == "acct-x"


def test_derive_rejects_access_token_without_exp(tmp_path: Path) -> None:
    sync = PiAuthSync(agent_dir=tmp_path / "pi", accounts_dir=tmp_path / "accounts")
    auth = tmp_path / "auth.json"
    auth.write_text(
        json.dumps(
            {
                "tokens": {
                    "access_token": _jwt({"no_exp": True}),
                    "refresh_token": "rt",
                    "account_id": "acct",
                }
            }
        ),
        encoding="utf-8",
    )

    with pytest.raises(PiAuthSyncError):
        sync._derive_credential(auth, "x")