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

ENV_PATH = "/root/site/yandex/.env_hram"
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["YANDEX2_CLIENT_ID"]
client_secret = creds["YANDEX2_CLIENT_SECRET"]
code = sys.argv[1].strip()

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 r:
        body = json.loads(r.read().decode())
except urllib.error.HTTPError as e:
    print("FAIL:", e.code, e.read().decode()[:300]); sys.exit(1)

if "access_token" not in body:
    print("NO TOKEN:", json.dumps(body, ensure_ascii=False)[:300]); sys.exit(1)

access = body["access_token"]
refresh = body.get("refresh_token", "")
expires_at = int(time.time()) + body.get("expires_in", 0)
print("OK: access len=%d, refresh=%s, expires_in=%d" % (len(access), "yes" if refresh else "no", body.get("expires_in", 0)))

def upsert(text, key, value):
    pat = re.compile(rf"^{re.escape(key)}=.*$", re.M)
    return pat.sub(f"{key}={value}", text) if pat.search(text) else text + f"\n{key}={value}\n"

content = open(ENV_PATH, encoding="utf-8").read()
for k, v in [("YANDEX2_ACCESS_TOKEN", access), ("YANDEX2_REFRESH_TOKEN", refresh),
             ("YANDEX2_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 via 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 r:
        d = json.loads(r.read().decode())
    print("VERIFY OK: user =", d.get("user", {}).get("display_name"),
          "| total = %.1f ГБ" % (d.get("total_space", 0) / 2**30),
          "| used = %.2f ГБ" % (d.get("used_space", 0) / 2**30))
except urllib.error.HTTPError as e:
    print("VERIFY FAIL:", e.code, e.read().decode()[:200])