import os, base64, time, requests
from pathlib import Path

ROOT = Path(r"C:\Users\Usuario\Downloads\SD")
BASE = ROOT / "_sandbox_cl_tagger" / "Uzaki Tsuki 104"
HARVEST = BASE / "harvest" / "por_imagen.txt"
OUT = ROOT / "Output" / "test_uzaki_fix"
OUT.mkdir(parents=True, exist_ok=True)

MASTER = ("masterpiece, best quality, highly detailed, sharp focus, very aesthetic, very awa, "
          "detailed background, score_9, score_8_up, <lora:IFL_v1.0_IL:0.7> <lora:add-detail-xl:0.5> "
          "<lora:blackwhiplash(bwl)-guy90-Illust-Lorav3:0.6> blaash, <lora:scara2b_ill_v1.00:0.7> scara2b-style, "
          "<lora:Add_color_ILXLv1:3>, 1girl, solo")
CHAR = "<lora:Tsuki_Uzaki:0.8> Tsuki Uzaki, grey_hair, low_tied_ponytail, closed_eyes, bangs, hair_between_eyes, huge_breasts, mature_female"
NEG = ("lowres, bad anatomy, bad hands, error, missing fingers, extra digit, fewer digits, cropped, worst quality, "
       "low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, bad feet, "
       "multiple views, splits, 2girls, extra people, duplicate_character, cloned_character, similar_character, "
       "ugly, deformed, disfigured, mutation, mutated, extra_limbs, extra_arms, extra_legs, fused_fingers, "
       "long_neck, cross-eyed, censored, mosaic censoring, bar censor, mirror, reflection, mirrored, glass_reflection, "
       "reflective_surface, 1boy, male, man, second person, multiple people, hetero, couple, extra person")

# parse harvest
entries = []
lines = HARVEST.read_text(encoding="utf-8").splitlines()
i = 0
while i < len(lines):
    if ".txt" in lines[i] and not lines[i].startswith("#"):
        img = lines[i].split()[0]
        tags = lines[i+1].strip() if i+1 < len(lines) else ""
        entries.append((img, tags))
        i += 2
    else:
        i += 1

# map image name -> tags
byimg = {e[0]: e[1] for e in entries}
wanted = ["1.txt","100.txt","102.txt","103.txt","104.txt","11.txt","17.txt"]

API = "http://127.0.0.1:7860/sdapi/v1/txt2img"
for idx, img in enumerate(wanted, 1):
    tags = byimg.get(img, "")
    # reinforce solo in block
    prompt = f"{MASTER} BREAK {CHAR}, half naked, solo BREAK {tags}, solo"
    payload = {
        "prompt": prompt,
        "negative_prompt": NEG,
        "width": 832,
        "height": 1216,
        "steps": 28,
        "sampler_name": "DPM++ 2M Karras",
        "cfg_scale": 7,
        "seed": 6000 + idx,
        "batch_size": 1,
        "n_iter": 1,
        "override_settings": {"CLIP_stop_at_last_layers": 2},
    }
    print(f"[{idx}/{len(wanted)}] {img} -> generating...")
    try:
        r = requests.post(API, json=payload, timeout=300)
        r.raise_for_status()
        data = r.json()
        b64 = data["images"][0]
        if "," in b64:
            b64 = b64.split(",", 1)[1]
        out_path = OUT / f"uzaki_fix_{idx:02d}_{img.replace('.txt','')}.png"
        out_path.write_bytes(base64.b64decode(b64))
        print(f"  OK -> {out_path.name}")
    except Exception as e:
        print(f"  ERROR: {e}")
    time.sleep(1)

print("DONE")