#!/usr/bin/env python3
"""Make Standard and Mini prefer small releases, and upgrade only toward small ones.

Run once per app; it is idempotent, so re-running updates everything in place.

- Adds "Small Encode" (release groups known for small HEVC/x264 encodes) and "720p" /
  "1080p" resolution formats. They score 0 in every profile except Standard and Mini.
- HD and UHD are otherwise untouched: they keep LQ at -10000, so they never pick these.
- Standard and Mini:
  - All allowed qualities share ONE group, so no "better resolution" or "better source"
    upgrade can ever download a big file. Only the custom-format score decides.
  - Scores: Small Encode +3000, x265 +2000 (so first grabs still prefer HEVC), LQ 0,
    release-tier formats 0 (one big release must never "beat" another big one).
    Standard adds +500 for 1080p and +250 for 720p; Mini adds +1000 for 720p and +500 for
    1080p. Those preferences stay below the upgrade threshold, so they never cause churn.
  - Upgrades are on, but an upgrade needs a score jump of at least 2500, and upgrading
    stops at 3000. In practice only gaining "Small Encode" clears that bar: a big H.264
    file gets replaced by a small encode, but never by a big HEVC one.

Usage: setup-small-profiles.py <sonarr|radarr>
"""
import json, re, subprocess, sys, urllib.request, pathlib

APP = sys.argv[1]
cfg = pathlib.Path.home() / f"Server/Media/volumes/{APP}/config.xml"
key = re.search(r"<ApiKey>([^<]+)", cfg.read_text()).group(1)
ip = subprocess.run(["docker", "inspect", APP, "--format", "{{range .NetworkSettings.Networks}}{{.IPAddress}} {{end}}"],
                    capture_output=True, text=True, check=True).stdout.split()[0]
base = f"http://{ip}:{8989 if APP == 'sonarr' else 7878}/api/v3"

def call(path, body=None, method=None):
    req = urllib.request.Request(base + path, method=method, headers={"X-Api-Key": key, "Content-Type": "application/json"},
                                 data=json.dumps(body).encode() if body is not None else None)
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.load(r)

GROUPS = r"^(MeGusta|PSA|ELiTE|GalaxyTV|GalaxyRG|GalaxyRG265|YIFY|YTS(\.(MX|LT|AG))?)$"
resolution = lambda r: [{"name": f"{r}p", "implementation": "ResolutionSpecification",
                         "negate": False, "required": True, "fields": [{"name": "value", "value": r}]}]
WANTED = {
    "Small Encode": [{"name": "Small encode groups", "implementation": "ReleaseGroupSpecification",
                      "negate": False, "required": False, "fields": [{"name": "value", "value": GROUPS}]}],
    "720p": resolution(720),
    "1080p": resolution(1080),
}
TIERS = {
    "Standard": {"Small Encode": 3000, "x265 (HD)": 2000, "LQ": 0, "LQ (Release Title)": 0, "1080p": 500, "720p": 250},
    "Mini":     {"Small Encode": 3000, "x265 (HD)": 2000, "LQ": 0, "LQ (Release Title)": 0, "720p": 1000, "1080p": 500},
}
RELEASE_TIER = re.compile(r"Tier \d+|WEB Scene")
MIN_UPGRADE_JUMP, UPGRADE_UNTIL = 2500, 3000

cfs = {c["name"]: c for c in call("/customformat")}
for name, specs in WANTED.items():
    body = {"name": name, "includeCustomFormatWhenRenaming": False, "specifications": specs}
    if name in cfs:
        body["id"] = cfs[name]["id"]; call(f"/customformat/{body['id']}", body, "PUT")
    else:
        call("/customformat", body)
cfs = {c["name"]: c["id"] for c in call("/customformat")}


def flatten(items):
    """Every quality in the profile with whether it is allowed, groups expanded, lowest rank first."""
    for it in items:
        for sub in (it["items"] or [it]):
            yield sub["quality"], it["allowed"]


def one_group(profile):
    """Put every allowed quality into a single top-ranked group; leave disallowed ones as they are."""
    flat = list(flatten(profile["items"]))
    group = {"id": 1100, "name": "Any allowed", "allowed": True,
             "items": [{"quality": q, "items": [], "allowed": True} for q, ok in flat if ok]}
    profile["items"] = [{"quality": q, "items": [], "allowed": False} for q, ok in flat if not ok] + [group]
    profile["cutoff"] = group["id"]
    return profile


for prof in call("/qualityprofile"):
    have = {f["format"] for f in prof["formatItems"]}
    for name in WANTED:
        if cfs[name] not in have:
            prof["formatItems"].append({"format": cfs[name], "name": name, "score": 0})
    scores = TIERS.get(prof["name"])
    if scores:
        prof = one_group(prof)
        prof.update(upgradeAllowed=True, minUpgradeFormatScore=MIN_UPGRADE_JUMP, cutoffFormatScore=UPGRADE_UNTIL)
        for f in prof["formatItems"]:
            if f["name"] in scores:
                f["score"] = scores[f["name"]]
            elif RELEASE_TIER.search(f["name"]):
                f["score"] = 0
    call(f"/qualityprofile/{prof['id']}", prof, "PUT")
    if scores:
        allowed = [q["quality"]["name"] for it in prof["items"] if it["allowed"] for q in it["items"]]
        print(f"{APP} {prof['name']}: one group of {len(allowed)} qualities ({allowed[0]} .. {allowed[-1]}); "
              f"upgrades on, jump >= {MIN_UPGRADE_JUMP}, until {UPGRADE_UNTIL}")
        print("   nonzero scores: " + ", ".join(f"{f['name']}={f['score']}" for f in prof["formatItems"] if f["score"]))