# -*- coding: utf-8 -*-
"""deploy_image.py — deploy gpt-image-2 (GlobalStandard) len 1 acc theo NHIEU region, TUAN TU.
Doc token/sub/proxy tu ket_qua/*/<slug>_session.json (khong login lai). Moi region:
  RG (1 cai/acc) -> Cognitive account (kind OpenAI, S0, customSubDomain) -> deployment gpt-image-2
  -> lay key1/key2 + endpoint -> ghi ket_qua/<date>/<slug>_gptimage2.txt.
Tao TUAN TU tung region + gap random (giam 715); gap 715 -> DUNG acc, giu region da xong.
KHONG dung foundry_flow (do hardcode Kimi/DataZoneStandard/feature Fireworks)."""
import os, sys, re, json, time, random, argparse, datetime, glob
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import armx

for _s in (sys.stdout, sys.stderr):
    try: _s.reconfigure(encoding="utf-8", errors="replace")
    except Exception: pass

V_CS = "2024-10-01"; V_KEYS = "2023-05-01"; V_RG = "2021-04-01"
MODEL = "gpt-image-2"; MVER = "2026-04-21"; MFMT = "OpenAI"
SKU = "GlobalStandard"; KIND = "OpenAI"
FILETAG = "gptimage2"      # hau to ten file txt ket qua (giu 'gptimage2' cho gpt-image-2 = backward-compat)
HERE = os.path.dirname(os.path.abspath(__file__))


def P(*a): print(*a, flush=True)


def _seg(s, n=48):
    s = re.sub(r"[^A-Za-z0-9-]", "-", (s or "")).strip("-")
    s = re.sub(r"-{2,}", "-", s)
    return s[:n].strip("-")


def _slug_of(f):
    return os.path.basename(f).replace("_session.json", "")


def _load_session(slug):
    best = None
    for f in glob.glob(os.path.join(HERE, "ket_qua", "*", "%s_session.json" % slug)):
        if best is None or os.path.getmtime(f) > os.path.getmtime(best):
            best = f
    if not best:
        raise SystemExit("khong co session cho %s" % slug)
    return json.load(open(best, encoding="utf-8"))


def _subid(s):
    v = s.get("sub") or ""
    return v.split("/")[2] if v.startswith("/subscriptions/") else v


def _is_715(blob):
    b = (blob or "").lower()
    return "715" in b or "unusual activity" in b


def create_rg(token, sub, rg, region, proxy):
    base = "/subscriptions/%s/resourcegroups/%s" % (sub, rg)
    d, st, _ = armx.arm_request(token, "GET", base, V_RG, None, proxy)
    if st == 200:
        return True, "co san"
    d, st, _ = armx.arm_request(token, "PUT", base, V_RG, {"location": region}, proxy)
    return (st < 400), ("tao moi" if st < 400 else armx._err(d))


def create_account(token, sub, rg, acct, region, proxy):
    apath = ("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.CognitiveServices/accounts/%s"
             % (sub, rg, acct))
    g, gs, _ = armx.arm_request(token, "GET", apath, V_CS, None, proxy)
    if gs == 200 and isinstance(g, dict) and (g.get("properties") or {}).get("provisioningState") == "Succeeded":
        return apath, "co san"
    body = {"location": region, "kind": KIND, "sku": {"name": "S0"},
            "identity": {"type": "SystemAssigned"},
            "properties": {"customSubDomainName": acct.lower(), "publicNetworkAccess": "Enabled"}}
    d, st, ok, why = armx.arm_lro(token, "PUT", apath, V_CS, body, proxy, poll_timeout=300, poll_interval=8)
    if not ok:
        return None, why
    return apath, "ok"


def deploy_model(token, apath, dn, cap, proxy):
    body = {"sku": {"name": SKU, "capacity": cap},
            "properties": {"model": {"format": MFMT, "name": MODEL, "version": MVER}}}
    dep = apath + "/deployments/%s" % dn
    last = ""
    for cp in (cap, 1):
        for attempt in range(5):
            d, st, ok, why = armx.arm_lro(token, "PUT", dep, V_CS,
                {"sku": {"name": SKU, "capacity": cp},
                 "properties": {"model": {"format": MFMT, "name": MODEL, "version": MVER}}},
                proxy, poll_timeout=300, poll_interval=8)
            if ok:
                return dep, "ok (cap=%d)" % cp
            last = why or ("st=%s" % st)
            blob = (str(why) + json.dumps(d) if isinstance(d, (dict, list)) else str(why)).lower()
            if _is_715(blob):
                return None, "715:" + last
            if st == 409 or "another operation" in blob or "parent" in blob:
                time.sleep(min(10 + attempt * 6, 40)); continue
            break   # loi khac -> thu cap=1 (vong ngoai) roi bo
    return None, last


def get_keys(token, sub, rg, acct, proxy):
    apath = ("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.CognitiveServices/accounts/%s"
             % (sub, rg, acct))
    k, _, _ = armx.arm_request(token, "POST", apath + "/listKeys", V_KEYS, None, proxy)
    a, _, _ = armx.arm_request(token, "GET", apath, V_CS, None, proxy)
    props = (a.get("properties") or {}) if isinstance(a, dict) else {}
    return {"key1": (k or {}).get("key1"), "key2": (k or {}).get("key2"),
            "endpoint": props.get("endpoint"),
            "openai": "https://%s.openai.azure.com/" % acct.lower()}


def run_region(token, sub, rg, slug, region, proxy, cap, dry):
    r = {"region": region, "ok": False, "error": None}
    acct = _seg("%s-i2-%s-%s" % (slug, region, ("%03x" % (random.randint(0, 4095)))), 44)
    dn = MODEL
    if dry:
        r["ok"] = True; r["acct"] = acct; r["note"] = "[dry] se tao acct=%s + deploy %s cap=%s" % (acct, MODEL, cap)
        return r
    ok_acc, why = create_account(token, sub, rg, acct, region, proxy)
    if not ok_acc:
        r["error"] = "account: " + str(why)[:160]; return r
    dep, dwhy = deploy_model(token, ok_acc, dn, cap, proxy)
    if not dep:
        # LUAT (user chot): TUYET DOI KHONG tu delete account/RG (delete co the kich 715).
        #   De account nguyen trang + note. Phong rac = QUOTA-FIRST o tang tren (chi deploy region co quota).
        r["error"] = "deploy: " + str(dwhy)[:160]; return r
    ke = get_keys(token, sub, rg, acct, proxy)
    r.update({"ok": True, "acct": acct, "deployment": dn, "keys": ke, "note": dwhy})
    return r


def write_txt(slug, sub, rg, results, skipped_715=None):
    date = datetime.date.today().isoformat()
    d = os.path.join(HERE, "ket_qua", date); os.makedirs(d, exist_ok=True)
    fp = os.path.join(d, "%s_%s.txt" % (slug, FILETAG))
    skipped_715 = skipped_715 or []
    okn = sum(1 for r in results if r.get("ok"))
    total = len(results) + len(skipped_715)
    hit715 = bool(skipped_715) or any(r.get("error") and _is_715(r.get("error")) for r in results)
    lines = ["# %s (%s) deploy — %s" % (MODEL, SKU, slug),
             "# sub=%s  rg=%s  ngay=%s" % (sub, rg, date),
             "# %d/%d region OK" % (okn, total)]
    if hit715:
        lines.append("# DUNG acc vi loi 715 (unusual activity). KHONG retry theo quy tac.")
    lines.append("")
    for r in results:
        if not r.get("ok"):
            tag = "FAIL 715" if _is_715(r.get("error")) else "FAIL"
            lines += ["[%s] %s: %s" % (r.get("region"), tag, r.get("error")), ""]; continue
        ke = r.get("keys") or {}
        lines += ["[%s] OK  %s" % (r.get("region"), r.get("note") or ""),
                  "  account    : %s" % r.get("acct"),
                  "  deployment : %s (%s v%s)" % (r.get("deployment"), MODEL, MVER),
                  "  endpoint   : %s" % ke.get("endpoint"),
                  "  openai_url : %s" % ke.get("openai"),
                  "  api_key    : %s" % ke.get("key1"),
                  "  api_key2   : %s" % ke.get("key2"), ""]
    for reg in skipped_715:
        lines += ["[%s] SKIP: khong tao (da dung acc vi 715)." % reg]
    if skipped_715:
        lines.append("")
    open(fp, "w", encoding="utf-8").write("\n".join(lines))
    return fp


def main():
    # global len dau: python yeu cau 'global' truoc MOI lan dung ten (ke ca doc default=MODEL o argparse)
    global MODEL, MFMT, MVER, SKU, KIND, FILETAG
    ap = argparse.ArgumentParser()
    ap.add_argument("--slug", required=True)
    ap.add_argument("--regions", required=True, help="csv, vd eastus2,uaenorth,swedencentral,westus3,polandcentral")
    ap.add_argument("--cap", type=int, default=2)
    ap.add_argument("--gap", default="20-45", help="giay nghi random giua region")
    ap.add_argument("--dry", action="store_true")
    # model-args: mac dinh gpt-image-2 (giu nguyen hanh vi cu). Truyen de deploy model khac.
    ap.add_argument("--model", default=MODEL, help="model.name (vd gpt-6-astra, claude-fable-5)")
    ap.add_argument("--format", dest="fmt", default=MFMT, help="model.format (OpenAI/Anthropic)")
    ap.add_argument("--version", default=MVER, help="model.version")
    ap.add_argument("--sku", default=SKU, help="GlobalStandard | DataZoneStandard")
    ap.add_argument("--kind", default=KIND, help="kind account CognitiveServices (OpenAI | AIServices)")
    args = ap.parse_args()

    # override hang so module tu args (cac ham deploy_model/create_account/write_txt doc global nay)
    MODEL, MFMT, MVER, SKU, KIND = args.model, args.fmt, args.version, args.sku, args.kind
    FILETAG = "gptimage2" if MODEL == "gpt-image-2" else re.sub(r"[^A-Za-z0-9]+", "-", MODEL).strip("-")

    s = _load_session(args.slug)
    token = s.get("token"); sub = _subid(s); proxy = s.get("proxy")
    if not token or not sub:
        raise SystemExit("session thieu token/sub")
    regions = [x.strip() for x in args.regions.split(",") if x.strip()]
    try:
        lo, hi = (float(x) for x in args.gap.split("-")) if "-" in args.gap else (float(args.gap) * 0.6, float(args.gap))
    except Exception:
        lo, hi = 20, 45

    rg = _seg("op-img2-%s-rg" % args.slug, 80)
    P("== deploy %s (%s) | slug=%s sub=%s | regions=%s | cap=%s dry=%s" % (MODEL, SKU, args.slug, sub, regions, args.cap, args.dry))
    if not args.dry:
        ok_rg, rgwhy = create_rg(token, sub, rg, regions[0], proxy)
        P("[RG] %s (%s)" % ("ok" if ok_rg else "LOI", rgwhy))
        if not ok_rg:
            raise SystemExit("RG loi: %s" % rgwhy)

    results = []; skipped_715 = []
    for i, reg in enumerate(regions):
        P("  -> [%d/%d] %s ..." % (i + 1, len(regions), reg))
        rr = run_region(token, sub, rg, args.slug, reg, proxy, args.cap, args.dry)
        results.append(rr)
        if rr.get("ok"):
            ke = rr.get("keys") or {}
            P("     OK acct=%s %s" % (rr.get("acct"), rr.get("note") or ""))
            if ke.get("endpoint"): P("        endpoint=%s" % ke.get("endpoint"))
        else:
            P("     FAIL: %s" % rr.get("error"))
            if _is_715(rr.get("error")):
                skipped_715 = regions[i + 1:]   # cac region con lai: KHONG tao, KHONG retry
                P("     !! 715 (unusual activity) -> DUNG acc nay, giu region da xong. KHONG retry.")
                if skipped_715:
                    P("     -> bo qua (715): %s" % ", ".join(skipped_715))
                break
        if not args.dry and i < len(regions) - 1:
            g = random.uniform(lo, hi); P("     [gap] nghi %.0fs ..." % g); time.sleep(g)

    fp = None
    if not args.dry:
        fp = write_txt(args.slug, sub, rg, results, skipped_715)
    okn = sum(1 for r in results if r.get("ok"))
    P("== XONG %s: %d/%d region OK%s" % (args.slug, okn, len(results), ("  file=%s" % fp) if fp else ""))
    print("__DONE__" + json.dumps({"slug": args.slug, "ok": okn, "total": len(results), "file": fp}, ensure_ascii=False))


if __name__ == "__main__":
    main()