from __future__ import annotations

import base64
import json
import unittest
from typing import Any
from unittest.mock import Mock, patch

import services.openai_backend_api as backend_module
from services.account_service import (
    AccountService,
    DEFAULT_CHATGPT_USER_AGENT,
    build_chatgpt_client_fingerprint,
    sanitize_registered_proxy,
)
from services.openai_backend_api import OpenAIBackendAPI


class MemoryStorage:
    def __init__(self, accounts: list[dict[str, Any]] | None = None) -> None:
        self.accounts = [dict(item) for item in (accounts or [])]

    def load_accounts(self) -> list[dict[str, Any]]:
        return [dict(item) for item in self.accounts]

    def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
        self.accounts = [dict(item) for item in accounts]


def make_jwt(payload: dict[str, Any]) -> str:
    def encode(value: dict[str, Any]) -> str:
        raw = json.dumps(value, separators=(",", ":")).encode("utf-8")
        return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")

    return f'{encode({"alg": "none"})}.{encode(payload)}.sig'


class FakeCookies:
    def __init__(self) -> None:
        self.values: list[tuple[str, str, str]] = []

    def set(self, name: str, value: str, *, domain: str) -> None:
        self.values.append((name, value, domain))


class FakeSession:
    def __init__(self) -> None:
        self.headers: dict[str, str] = {}
        self.cookies = FakeCookies()
        self.closed = False

    def close(self) -> None:
        self.closed = True


class AccountClientContextTests(unittest.TestCase):
    def test_account_context_is_migrated_and_stable_across_reload(self) -> None:
        access_token = make_jwt(
            {"https://api.openai.com/auth": {"chatgpt_account_id": "acct_stable"}}
        )
        storage = MemoryStorage([{"access_token": access_token}])

        first_service = AccountService(storage)
        first = first_service.get_account(access_token)
        second_service = AccountService(storage)
        second = second_service.get_account(access_token)

        self.assertEqual(first["account_id"], "acct_stable")
        self.assertTrue(first["account_key"])
        self.assertEqual(first["account_key"], second["account_key"])
        self.assertEqual(first["fp"], second["fp"])
        self.assertEqual(first["fp"]["impersonate"], "chrome145")
        self.assertIn("Chrome/145.", first["fp"]["user-agent"])

        rotated = build_chatgpt_client_fingerprint(
            "rotated-access-token",
            account_key=first["account_key"],
        )
        self.assertEqual(rotated["oai-device-id"], first["fp"]["oai-device-id"])
        self.assertEqual(rotated["oai-session-id"], first["fp"]["oai-session-id"])

    def test_legacy_mismatched_edge_profile_is_replaced_by_consistent_profile(self) -> None:
        fp = build_chatgpt_client_fingerprint(
            "token",
            {
                "user-agent": (
                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                    "Chrome/143.0.0.0 Safari/537.36 Edg/143.0.0.0"
                ),
                "impersonate": "edge101",
                "oai-device-id": "device-stable",
                "oai-session-id": "session-stable",
            },
        )

        self.assertEqual(fp["impersonate"], "chrome145")
        self.assertIn("Chrome/145.", fp["user-agent"])
        self.assertEqual(fp["oai-device-id"], "device-stable")
        self.assertEqual(fp["oai-session-id"], "session-stable")

    def test_backend_reuses_context_proxy_and_account_header_for_rotated_token(self) -> None:
        storage = MemoryStorage()
        service = AccountService(storage)
        fp = build_chatgpt_client_fingerprint(
            raw_fp={"user-agent": DEFAULT_CHATGPT_USER_AGENT, "impersonate": "chrome145"},
            device_id="login-device",
            session_id="login-session",
        )
        session = FakeSession()
        context = {
            "account_id": "acct_context",
            "fp": fp,
            "registered_proxy": "http://127.0.0.1:7890",
        }

        with (
            patch.object(backend_module, "account_service", service),
            patch.object(backend_module.requests, "Session", return_value=session) as factory,
        ):
            backend = OpenAIBackendAPI("rotated-token", account_context=context)

        self.assertEqual(
            factory.call_args.kwargs,
            {"impersonate": "chrome110", "verify": True, "proxy": "http://127.0.0.1:7890"},
        )
        self.assertEqual(backend.device_id, "login-device")
        self.assertEqual(backend.session_id, "login-session")
        self.assertEqual(session.headers["ChatGPT-Account-Id"], "acct_context")
        self.assertEqual(session.headers["OAI-Client-Version"], "prod-a194cd50d4416d3c0b47c740f206b12ce60f5887")
        self.assertEqual(session.headers["OAI-Client-Build-Number"], "6708908")
        self.assertEqual(session.headers["Sec-Ch-Ua-Platform-Version"], '"19.0.0"')
        self.assertEqual(backend._headers("/backend-api/me")["ChatGPT-Account-Id"], "acct_context")
        backend.close()
        self.assertTrue(session.closed)

    def test_discovered_account_id_is_applied_before_followup_requests_and_saved(self) -> None:
        storage = MemoryStorage([{"access_token": "opaque-token"}])
        service = AccountService(storage)
        session = FakeSession()

        with (
            patch.object(backend_module, "account_service", service),
            patch.object(backend_module.requests, "Session", return_value=session),
        ):
            backend = OpenAIBackendAPI("opaque-token")
            backend._get_default_account = Mock(return_value={"account_id": "acct_discovered", "plan_type": "free"})

            def get_me() -> dict[str, str]:
                self.assertEqual(session.headers["ChatGPT-Account-Id"], "acct_discovered")
                return {"email": "user@example.com", "id": "user-1"}

            backend._get_me = Mock(side_effect=get_me)
            backend._get_conversation_init = Mock(return_value={"limits_progress": []})
            result = backend.get_user_info()

        self.assertEqual(result["account_id"], "acct_discovered")
        self.assertEqual(service.get_account("opaque-token")["account_id"], "acct_discovered")

    def test_authenticated_proxy_is_not_persisted_in_account_data(self) -> None:
        self.assertEqual(sanitize_registered_proxy("http://user:secret@proxy.example:8080"), "")
        self.assertEqual(
            sanitize_registered_proxy("http://127.0.0.1:7890/path?token=secret"),
            "http://127.0.0.1:7890",
        )


if __name__ == "__main__":
    unittest.main()