#!/usr/bin/env python3
import json, re, urllib.request, urllib.parse, time, os

ENV_PATH = "/root/site/yandex/.env"

# --- read creds from .env ---
creds = {}
for line in open(ENV_PATH, encoding="utf-8"):
    line = line.strip()
    if line and not line.startswith("#") and "=" in line:
        k, v = line.split("=", 1)
        creds[k.strip()] = v.strip()

client_id = creds["YANDEX_CLIENT_ID"]
client_secret = creds["YANDEX_CLIENT_SECRET"]
code = "pxm7iuvo7vgmg5ue".strip()

# --- exchange code -> token ---
data = urllib.parse.urlencode({
    "grant_type": "authorization_code",
    "code": code,
    "client_id": client_id,
    "client_secret": client_secret,
}).encode()

req = urllib.request.Request(
    "https://oauth.yandex.ru/token",
    data=data,
    headers={"Content-Type": "application/x-www-form-urlencoded"},
)
try:
    with urllib.request.urlopen(req, timeout=30) as resp:
        body = json.loads(resp.read().decode())
except urllib.error.HTTPError as e:
    print("TOKEN EXCHANGE FAILED:", e.code, e.read().decode()[:500])
    raise SystemExit(1)

if "access_token" not in body:
    print("NO ACCESS TOKEN in response:", json.dumps(body, ensure_ascii=False)[:500])
    raise SystemExit(1)

access = body["access_token"]
refresh = body.get("refresh_token", "")
expires_in = body.get("expires_in", 0)
expires_at = int(time.time()) + expires_in

print("OK: got access_token (len=%d), refresh_token=%s, expires_in=%d" % (
    len(access), "yes" if refresh else "no", expires_in))

# --- save into .env (replace placeholder lines) ---
def upsert(text, key, value):
    pat = re.compile(rf"^{re.escape(key)}=.*$", re.M)
    if pat.search(text):
        return pat.sub(f"{key}={value}", text)
    return text + f"\n{key}={value}\n"

content = open(ENV_PATH, encoding="utf-8").read()
for k, v in [("YANDEX_ACCESS_TOKEN", access), ("YANDEX_REFRESH_TOKEN", refresh),
             ("YANDEX_TOKEN_EXPIRES_AT", str(expires_at))]:
    content = upsert(content, k, v)
open(ENV_PATH, "w", encoding="utf-8").write(content)
os.chmod(ENV_PATH, 0o600)
print("Saved to", ENV_PATH)

# --- verify: Disk API ---
try:
    req = urllib.request.Request(
        "https://cloud-api.yandex.net/v1/disk",
        headers={"Authorization": f"OAuth {access}"},
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        d = json.loads(resp.read().decode())
    print("VERIFY Disk API OK: user =", d.get("user", {}).get("display_name"),
          "| total space =", d.get("total_space"), "| used =", d.get("used_space"))
except urllib.error.HTTPError as e:
    print("VERIFY Disk API FAILED:", e.code, e.read().decode()[:300])