#!/usr/bin/env python3
"""
A2A Consumer Rotator (HF VM Pull Script)
Pulls a fresh OpenAI refresh token from SIN-Supabase and writes the
OpenCode auth schema to ~/.local/share/opencode/auth.json.
"""

from __future__ import annotations

import json
import os
import sys
from collections.abc import Mapping
from pathlib import Path

import requests


SUPABASE_URL = os.environ.get("SUPABASE_URL", "").strip()
SUPABASE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "").strip()
AUTH_FILE = Path.home() / ".local" / "share" / "opencode" / "auth.json"

JsonScalar = str | int | float | bool | None
JsonValue = JsonScalar | dict[str, "JsonValue"] | list["JsonValue"]


def extract_refresh_token(row: Mapping[str, JsonValue]) -> str | None:
    direct_token = row.get("refresh_token") or row.get("token")
    if isinstance(direct_token, str) and direct_token.strip():
        return direct_token.strip()

    auth_json_data = row.get("auth_json_data")
    if isinstance(auth_json_data, Mapping):
        openai_auth = auth_json_data.get("openai")
        if isinstance(openai_auth, Mapping):
            refresh_token = openai_auth.get("refresh")
            if isinstance(refresh_token, str) and refresh_token.strip():
                return refresh_token.strip()

        refresh_token = auth_json_data.get("refresh") or auth_json_data.get("refresh_token")
        if isinstance(refresh_token, str) and refresh_token.strip():
            return refresh_token.strip()

    return None


def build_headers() -> dict[str, str]:
    return {
        "apikey": SUPABASE_KEY,
        "Authorization": f"Bearer {SUPABASE_KEY}",
        "Content-Type": "application/json",
        "Prefer": "return=representation",
    }


def pull_fresh_token() -> bool:
    if not SUPABASE_URL or not SUPABASE_KEY:
        print("ERROR: SUPABASE_URL or SUPABASE_SERVICE_ROLE_KEY is missing.")
        return False

    base_url = SUPABASE_URL.rstrip("/")
    headers = build_headers()

    print("Looking for a fresh token in Supabase...")
    get_url = (
        f"{base_url}/rest/v1/openai_account_pool"
        "?status=eq.FRESH&select=id,refresh_token,token,auth_json_data&limit=1"
    )
    response = requests.get(get_url, headers=headers, timeout=10)
    response.raise_for_status()

    payload = response.json()
    if not isinstance(payload, list):
        print("ERROR: Supabase returned an unexpected payload shape.")
        return False

    rows = [row for row in payload if isinstance(row, dict)]
    if not rows:
        print("ERROR: No fresh token available in openai_account_pool.")
        return False

    row = rows[0]
    row_id_value = row.get("id")
    if not isinstance(row_id_value, (str, int)):
        print("ERROR: Supabase row is missing a valid id.")
        return False

    row_id = str(row_id_value)
    refresh_token = extract_refresh_token(row)
    if not refresh_token:
        print(f"ERROR: Row {row_id} does not contain a usable refresh token.")
        return False

    patch_url = f"{base_url}/rest/v1/openai_account_pool?id=eq.{row_id}"
    patch_response = requests.patch(
        patch_url,
        headers=headers,
        json={"status": "IN_USE"},
        timeout=10,
    )
    patch_response.raise_for_status()

    auth_payload = {
        "openai": {
            "type": "oauth",
            "refresh": refresh_token,
        }
    }
    _ = AUTH_FILE.parent.mkdir(parents=True, exist_ok=True)
    _ = AUTH_FILE.write_text(json.dumps(auth_payload, indent=2) + "\n", encoding="utf-8")
    print(f"Updated {AUTH_FILE} from Supabase row {row_id}.")
    return True


if __name__ == "__main__":
    sys.exit(0 if pull_fresh_token() else 1)