# -*- coding: utf-8 -*-
import os, io, re, random, string, tempfile
import spaces
import gradio as gr
from PIL import Image

# --- CONSTANTS ---
POTRACE_BIN = 'potrace'
UPM = 1000
CROP = 8
SIDEBEARING = 20
GLYPH_SCALE = 0.112
PRECISION = 1
HF_TOKEN = os.environ.get("HF_TOKEN", "")
DEFAULT_IMG_PROMPT = "Design a custom typeface that takes direct inspiration on the attached control image. The font should faithfully replicate the unique style of the reference."

# --- 6x6 MAPPINGS ---
MAP_UC = [['A','B','C','D','E','F'],['G','H','I','J','K','L'],['M','N','O','P','Q','R'],['S','T','U','V','W','X'],['Y','Z','Æ','Œ','Ø','ß'],['Ç','.',',',';','?','!']]
MAP_LC = [['a','b','c','d','e','f'],['g','h','i','j','k','l'],['m','n','o','p','q','r'],['s','t','u','v','w','x'],['y','z','æ','œ','ø','Ð'],['ç',':',"'",'"','«','<']]
MAP_PUNC = [['1','2','3','4','5','6'],['7','8','9','0','#','%'],['(','[','{','$','€','£'],['&','@','_','+','-','='],['*','/','^','°','→','—'],['`','´','ˆ','¨','˜','•']]

# --- PER-GRID BASELINE REFERENCES ---
BASELINE_REFS = {
    0: ['H', 'I', 'E', 'A', 'B', 'D', 'F', 'L', 'M', 'N', 'P', 'R', 'T', 'U', 'V', 'X', 'Y', 'Z'],
    1: ['n', 'm', 'u', 'x', 'h', 'i', 'l', 'k', 'r', 'v', 'w', 'z', 'a', 'e', 'o'],
    2: ['1', '0', '2', '4', '7', '8', '9', '#', '%', '$', '£', '&', '@', '+', '='],
}

# --- SVG PROCESSING ---
def simplify_svg_path(d):
    from fontTools.pens.recordingPen import RecordingPen
    from fontTools.pens.svgPathPen import SVGPathPen
    from fontTools.svgLib.path import parse_path
    rec = RecordingPen()
    parse_path(d, rec)
    svg = SVGPathPen(None)
    rec.replay(svg)
    return re.sub(r"(\d+\.\d+)", lambda m: f"{float(m.group(1)):.{PRECISION}f}", svg.getCommands())

def get_char_path_rdp(pil_img, upscale=2):
    import tempfile, subprocess
    import cv2, numpy as np
    rgb = np.array(pil_img.convert("RGB"))
    black_mask = (
        (rgb[:,:,0] < 80) &
        (rgb[:,:,1] < 80) &
        (rgb[:,:,2] < 80)
    ).astype(np.uint8) * 255
    h, w = black_mask.shape
    up = cv2.resize(black_mask, (w*upscale, h*upscale), interpolation=cv2.INTER_NEAREST)
    with tempfile.NamedTemporaryFile(suffix='.pbm', delete=False) as f:
        pbm = f.name
    with tempfile.NamedTemporaryFile(suffix='.svg', delete=False) as f:
        svg = f.name
    Image.fromarray(cv2.flip(255-up, 0)).convert('1').save(pbm)
    try:
        subprocess.run([POTRACE_BIN, pbm, '-s', '-o', svg, '--alphamax', '0.9', '--opttolerance', '0.5', '--turdsize', '30'], check=True, capture_output=True)
    except:
        return ""
    finally:
        if os.path.exists(pbm):
            os.unlink(pbm)
    if not os.path.exists(svg):
        return ""
    with open(svg) as f:
        content = f.read()
    os.unlink(svg)
    paths = re.findall(r'd="([^"]+)"', content)
    return ' '.join(simplify_svg_path(p) for p in paths).strip() if paths else ""

# --- OTF CONSTRUCTION ---
MIRROR_MAP = {')': '(', ']': '[', '}': '{'}

ACCENT_MAP = {
    'à':('a','`'),'è':('e','`'),'ì':('i','`'),'ò':('o','`'),'ù':('u','`'),
    'À':('A','`'),'È':('E','`'),'Ì':('I','`'),'Ò':('O','`'),'Ù':('U','`'),
    'á':('a','´'),'é':('e','´'),'í':('i','´'),'ó':('o','´'),'ú':('u','´'),
    'Á':('A','´'),'É':('E','´'),'Í':('I','´'),'Ó':('O','´'),'Ú':('U','´'),
    'â':('a','ˆ'),'ê':('e','ˆ'),'î':('i','ˆ'),'ô':('o','ˆ'),'û':('u','ˆ'),
    'Â':('A','ˆ'),'Ê':('E','ˆ'),'Î':('I','ˆ'),'Ô':('O','ˆ'),'Û':('U','ˆ'),
    'ä':('a','¨'),'ë':('e','¨'),'ï':('i','¨'),'ö':('o','¨'),'ü':('u','¨'),
    'Ä':('A','¨'),'Ë':('E','¨'),'Ï':('I','¨'),'Ö':('O','¨'),'Ü':('U','¨'),
    'ã':('a','˜'),'õ':('o','˜'),'ñ':('n','˜'),
    'Ã':('A','˜'),'Õ':('O','˜'),'Ñ':('N','˜'),
}

def build_otf(images, font_name):
    from fontTools.fontBuilder import FontBuilder
    from fontTools.pens.t2CharStringPen import T2CharStringPen
    from fontTools.pens.boundsPen import BoundsPen
    from fontTools.pens.transformPen import TransformPen
    from fontTools.pens.recordingPen import RecordingPen
    from fontTools.misc.transform import Identity
    from fontTools.svgLib.path import parse_path
    from fontTools.cffLib import PrivateDict
    from fontTools.ttLib import newTable
    from fontTools.ttLib.tables import _k_e_r_n as _kern
    import tempfile, subprocess, cv2, numpy as np
    from PIL import Image, ImageDraw

    all_maps = [MAP_UC, MAP_LC, MAP_PUNC]
    glyph_data = {}

    for idx, img in enumerate(images):
        w, h = img.size
        arr = np.array(img.convert("RGB"))
        arr_up = cv2.resize(arr, (w*2, h*2), interpolation=cv2.INTER_CUBIC)
        img_up = Image.fromarray(arr_up)
        w2, h2 = img_up.size
        cw, ch = w2//6, h2//6
        scale = 800/ch
        current_map = all_maps[idx]
        for r in range(6):
            for c in range(6):
                char = current_map[r][c]
                box = (c*cw, r*ch, (c+1)*cw, (r+1)*ch)
                d = get_char_path_rdp(img_up.crop(box), upscale=1)
                if not d:
                    continue
                bp = BoundsPen(None)
                try:
                    parse_path(d, bp)
                except:
                    continue
                if bp.bounds:
                    glyph_data[char] = {'d': d, 'b': bp.bounds, 'scale': scale, 'grid': idx}

    # Computed baseline
    baseline_per_grid = {}
    for grid_idx in range(len(all_maps)):
        chosen_ref = None
        for ref_char in BASELINE_REFS.get(grid_idx, []):
            if ref_char in glyph_data and glyph_data[ref_char].get('grid') == grid_idx:
                gd = glyph_data[ref_char]
                baseline_per_grid[grid_idx] = gd['b'][3] * gd['scale'] * GLYPH_SCALE
                chosen_ref = ref_char
                break
        if chosen_ref is None:
            baseline_per_grid[grid_idx] = baseline_per_grid.get(0, 0)
            print(f"⚠️  Grid {grid_idx}: no reference glyph found, fallback = {baseline_per_grid[grid_idx]:.1f}")
        else:
            print(f"📐 Grid {grid_idx}: baseline = {baseline_per_grid[grid_idx]:.1f} (ref '{chosen_ref}')")

    def get_baseline(char):
        if char in glyph_data:
            return baseline_per_grid.get(glyph_data[char]['grid'], baseline_per_grid.get(0, 0))
        return baseline_per_grid.get(0, 0)

    dummy = PrivateDict()
    dummy.nominalWidthX = 0
    cs = {".notdef": T2CharStringPen(600, None).getCharString(private=dummy)}
    metrics = {".notdef": (600, 0)}
    cmap, order = {}, [".notdef"]
    cs["space"] = T2CharStringPen(300, None).getCharString(private=dummy)
    metrics["space"] = (300, 0)
    order.append("space")
    cmap[32] = "space"

    font_recordings = {}
    font_bounds = {}
    font_widths = {}
    font_centers = {}

    for char, data in glyph_data.items():
        if len(char) != 1:
            continue
        b, s = data['b'], data['scale']
        width = round((b[2]-b[0]) * s * GLYPH_SCALE + SIDEBEARING * 2)
        tx = SIDEBEARING - b[0] * s * GLYPH_SCALE
        ty = get_baseline(char)
        transform = Identity.translate(tx, ty).scale(s * GLYPH_SCALE, -s * GLYPH_SCALE)
        rec = RecordingPen()
        parse_path(data['d'], TransformPen(rec, transform))
        font_recordings[char] = rec.value
        font_widths[char] = width
        bp2 = BoundsPen(None)
        rec2 = RecordingPen()
        rec2.value = rec.value
        rec2.replay(bp2)
        bounds = bp2.bounds or (0, 0, width, 800)
        font_bounds[char] = bounds
        font_centers[char] = ((bounds[0] + bounds[2]) / 2, (bounds[1] + bounds[3]) / 2)
        pen = T2CharStringPen(width, None)
        rec3 = RecordingPen()
        rec3.value = rec.value
        rec3.replay(pen)
        glyph_name = f"glyph{ord(char)}"
        cs[glyph_name] = pen.getCharString(private=dummy)
        order.append(glyph_name)
        cmap[ord(char)] = glyph_name
        metrics[glyph_name] = (width, 0)

    # Mirrors
    for dst_char, src_char in MIRROR_MAP.items():
        if src_char not in font_recordings:
            continue
        w = font_widths[src_char]
        pen = T2CharStringPen(w, None)
        flip = TransformPen(pen, (-1, 0, 0, 1, w, 0))
        rec = RecordingPen()
        rec.value = font_recordings[src_char]
        rec.replay(flip)
        gn = f"glyph{ord(dst_char)}"
        cs[gn] = pen.getCharString(private=dummy)
        order.append(gn)
        cmap[ord(dst_char)] = gn
        metrics[gn] = (w, 0)
        if src_char in font_centers:
            cx, cy = font_centers[src_char]
            font_centers[dst_char] = (w - cx, cy)

    # Accented
    for dst_char, (base_char, accent_char) in ACCENT_MAP.items():
        if base_char not in font_recordings or accent_char not in font_recordings:
            continue
        w = font_widths[base_char]
        bb = font_bounds[base_char]
        ab = font_bounds[accent_char]
        dx = ((bb[0]+bb[2]) - (ab[0]+ab[2])) / 2
        dy = bb[3] - ab[1] + 40
        pen = T2CharStringPen(w, None)
        rec1 = RecordingPen()
        rec1.value = font_recordings[base_char]
        rec1.replay(pen)
        rec2 = RecordingPen()
        rec2.value = font_recordings[accent_char]
        rec2.replay(TransformPen(pen, (1, 0, 0, 1, dx, dy)))
        gn = f"glyph{ord(dst_char)}"
        cs[gn] = pen.getCharString(private=dummy)
        order.append(gn)
        cmap[ord(dst_char)] = gn
        metrics[gn] = (w, 0)
        if base_char in font_centers and accent_char in font_centers:
            bx, by = font_centers[base_char]
            ax, ay = font_centers[accent_char]
            font_centers[dst_char] = (bx + dx, by + dy)

    # Bracket mirrors
    MIRROR_PAIRS = {'(': ')', '[': ']', '{': '}'}
    for src_char, dst_char in MIRROR_PAIRS.items():
        if src_char not in glyph_data:
            continue
        data = glyph_data[src_char]
        b, s = data['b'], data['scale']
        width = round((b[2]-b[0]) * s * GLYPH_SCALE + SIDEBEARING * 2)
        pen = T2CharStringPen(width, None)
        tx = SIDEBEARING + b[2] * s * GLYPH_SCALE
        ty = get_baseline(src_char)
        transform = Identity.translate(tx, ty).scale(-s * GLYPH_SCALE, -s * GLYPH_SCALE)
        parse_path(data['d'], TransformPen(pen, transform))
        glyph_name = f"glyph{ord(dst_char)}"
        if glyph_name not in cs:
            cs[glyph_name] = pen.getCharString(private=dummy)
            order.append(glyph_name)
            cmap[ord(dst_char)] = glyph_name
            metrics[glyph_name] = (width, 0)
            if src_char in font_centers:
                cx, cy = font_centers[src_char]
                font_centers[dst_char] = (width - cx, cy)

    # Combining-accent aliases
    ACCENT_ALIASES = {
        '`': [0x0060, 0x0300], '´': [0x00B4, 0x0301],
        'ˆ': [0x02C6, 0x0302], '¨': [0x00A8, 0x0308], '˜': [0x02DC, 0x0303],
    }
    for base_char, codepoints in ACCENT_ALIASES.items():
        if base_char not in glyph_data:
            continue
        existing_glyph_name = f"glyph{ord(base_char)}"
        if existing_glyph_name in cs:
            for cp in codepoints[1:]:
                if cp not in cmap:
                    cmap[cp] = existing_glyph_name

    fb = FontBuilder(UPM, isTTF=False)
    fb.setupGlyphOrder(order)
    fb.setupCharacterMap(cmap)
    fb.setupCFF(
        font_name,
        {"FullName": font_name, "FamilyName": font_name, "Weight": "Regular"},
        cs,
        {font_name: dummy}
    )
    fb.setupHorizontalMetrics(metrics)
    fb.setupHorizontalHeader(ascent=REF["asc"], descent=REF["dsc"])
    fb.setupNameTable({
        "familyName": font_name,
        "styleName": "Regular",
        "uniqueFontIdentifier": f"{font_name}:Version 1.000",
        "fullName": font_name,
        "version": "Version 1.000",
        "psName": font_name
    })
    fb.setupOS2(
        sTypoAscender=REF["tAsc"],
        sTypoDescender=REF["tDsc"],
        sTypoLineGap=REF["tGap"],
        usWinAscent=REF["wAsc"],
        usWinDescent=REF["wDsc"],
        sxHeight=REF["xH"],
        sCapHeight=REF["cH"],
        fsType=4,
        fsSelection=64
    )
    fb.setupPost(italicAngle=0, underlinePosition=REF["ulP"], underlineThickness=REF["ulT"])

    # -------------------------------------------------------------------------
    # BUBBLE KERNING — pure vector, written to GPOS (no pair-count limit)
    # -------------------------------------------------------------------------
    try:
        import numpy as np
        from fontTools.otlLib.builder import buildValue, buildPairPosGlyphs
        from fontTools.ttLib import newTable
        from fontTools.ttLib.tables import otTables

        # --- Parameters ---
        BUBBLE_RADIUS  = 10   # bubble radius in UPM, must be < SIDEBEARING (20)
        N_SAMPLES      = 300  # sampled points per contour
        KERN_THRESHOLD = -2   # pairs with kern >= threshold are skipped (sub-UPM noise)
        KERN_CAP       = -280 # anti-collision floor

        def sample_contour(recording, n=N_SAMPLES):
            pts = []

            def add_line(p0, p1):
                d = ((p1[0]-p0[0])**2 + (p1[1]-p0[1])**2) ** 0.5
                steps = max(1, int(d / 8))
                for i in range(steps + 1):
                    t = i / steps
                    pts.append((p0[0] + t*(p1[0]-p0[0]),
                                p0[1] + t*(p1[1]-p0[1])))

            def add_cubic(p0, p1, p2, p3):
                d = ((p3[0]-p0[0])**2 + (p3[1]-p0[1])**2) ** 0.5
                steps = max(4, int(d * 1.5 / 8))
                for i in range(steps + 1):
                    t  = i / steps
                    mt = 1.0 - t
                    bx = mt**3*p0[0] + 3*mt**2*t*p1[0] + 3*mt*t**2*p2[0] + t**3*p3[0]
                    by = mt**3*p0[1] + 3*mt**2*t*p1[1] + 3*mt*t**2*p2[1] + t**3*p3[1]
                    pts.append((bx, by))

            cur = (0.0, 0.0)
            start = (0.0, 0.0)
            for op, args in recording:
                if op == 'moveTo':
                    cur = args[0]; start = cur
                elif op == 'lineTo':
                    add_line(cur, args[0]); cur = args[0]
                elif op == 'curveTo':
                    add_cubic(cur, args[0], args[1], args[2]); cur = args[2]
                elif op in ('endPath', 'closePath'):
                    if cur != start:
                        add_line(cur, start)
                    cur = start

            if not pts:
                return None
            arr = np.array(pts, dtype=np.float32)
            if len(arr) > n:
                idx = np.round(np.linspace(0, len(arr)-1, n)).astype(int)
                arr = arr[idx]
            return arr

        # Sampling
        chars_available = [c for c in font_recordings if len(c) == 1]
        contour_points  = {}
        for char in chars_available:
            pts = sample_contour(font_recordings[char])
            if pts is not None and len(pts) >= 2:
                contour_points[char] = pts

        print(f"📐 Contours sampled: {len(contour_points)} glyphs")

        # All-to-all computation
        diameter   = BUBBLE_RADIUS * 2.0
        gpos_pairs = {}   # (glyph_name_L, glyph_name_R) -> kern_upm

        for l_char in contour_points:
            pts_L = contour_points[l_char]
            adv_L = float(font_widths[l_char])
            lg    = f"glyph{ord(l_char)}"
            if lg not in metrics:
                continue

            for r_char in contour_points:
                rg = f"glyph{ord(r_char)}"
                if rg not in metrics:
                    continue

                pts_R_shifted = contour_points[r_char] + np.array([adv_L, 0.0], dtype=np.float32)
                diff     = pts_L[:, np.newaxis, :] - pts_R_shifted[np.newaxis, :, :]
                min_dist = float(np.sqrt((diff**2).sum(axis=2)).min())
                kern_upm = int(round(diameter - min_dist))

                if l_char in ('A','V','W') and r_char in ('A','V','W'):
                    print(f"  {l_char}->{r_char} : min_dist={min_dist:.1f}  kern={kern_upm}")

                if kern_upm < KERN_THRESHOLD:
                    gpos_pairs[(lg, rg)] = max(kern_upm, KERN_CAP)

        print(f"✅ Bubble Kerning (vector): {len(gpos_pairs)} pairs")

        if gpos_pairs:
            val0 = buildValue({})
            pairs_for_builder = {
                (lg, rg): (buildValue({"XAdvance": kern}), val0)
                for (lg, rg), kern in gpos_pairs.items()
            }

            glyph_map = fb.font.getReverseGlyphMap()
            subtables = buildPairPosGlyphs(pairs_for_builder, glyph_map)

            lookup               = otTables.Lookup()
            lookup.LookupType    = 2
            lookup.LookupFlag    = 0
            lookup.SubTableCount = len(subtables)
            lookup.SubTable      = subtables

            lookup_list               = otTables.LookupList()
            lookup_list.Lookup        = [lookup]
            lookup_list.LookupCount   = 1

            feature_record            = otTables.FeatureRecord()
            feature_record.FeatureTag = "kern"
            feature                   = otTables.Feature()
            feature.LookupListIndex   = [0]
            feature.LookupCount       = 1
            feature_record.Feature    = feature

            lang_sys              = otTables.DefaultLangSys()
            lang_sys.ReqFeatureIndex = 0xFFFF
            lang_sys.FeatureIndex = [0]
            lang_sys.FeatureCount = 1
            lang_sys.LookupOrderOffset = 0

            script               = otTables.Script()
            script.DefaultLangSys = lang_sys
            script.LangSysCount  = 0
            script.LangSysRecord = []

            script_record            = otTables.ScriptRecord()
            script_record.ScriptTag  = "DFLT"
            script_record.Script     = script

            script_list              = otTables.ScriptList()
            script_list.ScriptRecord = [script_record]
            script_list.ScriptCount  = 1

            feature_list                 = otTables.FeatureList()
            feature_list.FeatureRecord   = [feature_record]
            feature_list.FeatureCount    = 1

            gpos_table          = otTables.GPOS()
            gpos_table.Version  = 0x00010000
            gpos_table.ScriptList  = script_list
            gpos_table.FeatureList = feature_list
            gpos_table.LookupList  = lookup_list

            gpos                = newTable("GPOS")
            gpos.table          = gpos_table
            fb.font["GPOS"]     = gpos
            print(f"✅ GPOS written: {len(gpos_pairs)} kern pairs")

    except Exception as e:
        print(f"⚠️ Bubble Kerning failed: {e}")
        import traceback
        traceback.print_exc()

    # --- Final OTF save ---
    tmp_otf = f"/tmp/{font_name}.otf"
    fb.save(tmp_otf)
    with open(tmp_otf, "rb") as f:
        data = f.read()
    os.remove(tmp_otf)
    return data

# --- REFERENCE METRICS (Helvetica LT Std Regular) ---
REF = {
    "asc": 718, "dsc": -282, "tAsc": 718, "tDsc": -282, "tGap": 200,
    "wAsc": 931, "wDsc": 225, "xH": 524, "cH": 718, "ulP": -75, "ulT": 50
}

# =============================================================================
# MODEL LOADING (module scope, once at Space startup)
# ZeroGPU: .to("cuda") here is fine, the real CUDA init is deferred by `spaces`.
# =============================================================================
import torch
from diffusers import Flux2KleinPipeline
from huggingface_hub import login

if HF_TOKEN:
    login(token=HF_TOKEN)

print("🚀 Loading Flux.2-klein + LoRAs ...")
pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-base-4B",
    torch_dtype=torch.bfloat16,
)
pipe.load_lora_weights("ChevalierJoseph/TYPOTOPIA_APP", weight_name="typotopiaMAJ.safetensors", adapter_name="uc")
pipe.load_lora_weights("ChevalierJoseph/TYPOTOPIA_APP", weight_name="typotopiaMIN.safetensors", adapter_name="lc")
pipe.load_lora_weights("ChevalierJoseph/TYPOTOPIA_APP", weight_name="typotopiaPONCT.safetensors", adapter_name="punc")
pipe.to("cuda")
print("✅ Model ready!")


# =============================================================================
# GPU INFERENCE — only the 3 FLUX passes run under @spaces.GPU
# duration = max GPU allocation per request (counts against the ZeroGPU quota).
# =============================================================================
@spaces.GPU(duration=180)
def run_pipeline(enhanced, input_img, seed):
    gen = torch.Generator("cuda").manual_seed(int(seed))

    pipe.set_adapters("uc")
    img1 = pipe(
        prompt=f"[typotopiaMAJ], {enhanced}",
        image=input_img,
        num_inference_steps=5,
        guidance_scale=7 if input_img else 3.5,
        height=1536, width=1536,
        generator=gen,
    ).images[0]

    pipe.set_adapters("lc")
    img2 = pipe(
        prompt="[typotopiaMIN], Design a custom lowercase typeface inspired by the attached control image.",
        image=img1,
        num_inference_steps=5,
        guidance_scale=7,
        height=1536, width=1536,
        generator=gen,
    ).images[0]

    pipe.set_adapters("punc")
    img3 = pipe(
        prompt="[typotopiaPONCT], Design a custom punctuation typeface inspired by the attached control image.",
        image=img1,
        num_inference_steps=5,
        guidance_scale=7,
        height=1536, width=1536,
        generator=gen,
    ).images[0]

    return img1, img2, img3


# =============================================================================
# ORCHESTRATOR + GRADIO UI
# =============================================================================
def _render_preview(otf_path, font_name):
    try:
        from PIL import Image, ImageDraw, ImageFont
        sample = f"{font_name}\nAa Bb Cc 123\nThe quick brown fox"
        img = Image.new("RGB", (1700, 620), "white")
        draw = ImageDraw.Draw(img)
        font = ImageFont.truetype(otf_path, 150)
        draw.multiline_text((40, 30), sample, font=font, fill="black", spacing=30)
        return img
    except Exception as e:
        print(f"⚠️ Preview failed: {e}")
        return None


def generate(prompt, control_image, progress=gr.Progress(track_tqdm=True)):
    seed = random.randint(0, 2**32 - 1)

    input_img = None
    prefixes = ["Vex","Aur","Kyr","Lux","Nox","Arc","Sol","Vel","Fen","Zor","Cal","Dex","Ora","Pax","Rux"]
    suffixes = ["ra","is","on","us","ia","el","an","ix","em","or","al","en","ax","um","yr"]
    font_name = random.choice(prefixes) + random.choice(suffixes)

    if control_image is not None:
        input_img = control_image.convert("RGB").resize((1536, 1536))
        enhanced = DEFAULT_IMG_PROMPT
    else:
        if not (prompt and prompt.strip()):
            raise gr.Error("Provide a prompt OR a control image.")
        enhanced = prompt  # raw user input, no refinement

    # --- GPU: 3 FLUX passes ---
    img1, img2, img3 = run_pipeline(enhanced, input_img, seed)

    if not font_name:
        font_name = "Font" + "".join(random.choices(string.ascii_uppercase, k=3))

    # --- CPU: potrace vectorization + OTF build ---
    otf_bytes = build_otf([img1, img2, img3], font_name)

    out_dir = tempfile.mkdtemp()
    otf_path = os.path.join(out_dir, f"{font_name}.otf")
    with open(otf_path, "wb") as f:
        f.write(otf_bytes)

    gallery = [(img1, "Uppercase"), (img2, "Lowercase"), (img3, "Punctuation")]
    preview = _render_preview(otf_path, font_name)
    return gallery, preview, otf_path


# =============================================================================
# OPTIONAL CUSTOM TITLE FONT
# Drop your own font in the repo as one of these names and the hero title will
# use it (otherwise it falls back to the Helvetica/Archivo grotesque stack):
#   title.otf / title.ttf / title.woff2 / title.woff   (or under assets/)
# =============================================================================
def _load_title_font():
    import base64
    fmt  = {"otf": "opentype", "ttf": "truetype", "woff": "woff", "woff2": "woff2"}
    mime = {"otf": "font/otf", "ttf": "font/ttf", "woff": "font/woff", "woff2": "font/woff2"}
    candidates = [
        "title.otf", "title.ttf", "title.woff2", "title.woff",
        "assets/title.otf", "assets/title.ttf", "assets/title.woff2", "assets/title.woff",
    ]
    for path in candidates:
        if os.path.exists(path):
            ext = path.rsplit(".", 1)[-1].lower()
            with open(path, "rb") as fh:
                b64 = base64.b64encode(fh.read()).decode()
            print(f"🔤 Title font: {path}")
            face = (
                "@font-face {"
                "font-family: 'TitleFont';"
                f"src: url(data:{mime[ext]};base64,{b64}) format('{fmt[ext]}');"
                "font-weight: 100 900; font-style: normal; font-display: swap;}"
            )
            return face, "'TitleFont', 'Helvetica Neue', Helvetica, 'Archivo', Arial, sans-serif"
    return "", "'Helvetica Neue', Helvetica, 'Archivo', Arial, sans-serif"

TITLE_FONT_FACE, TITLE_FAMILY = _load_title_font()

# =============================================================================
# THEME — International Typographic Style (light grey / black / violet)
# =============================================================================
THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.violet,
    neutral_hue=gr.themes.colors.gray,
    radius_size=gr.themes.sizes.radius_none,
    font=["Helvetica Neue", "Helvetica", gr.themes.GoogleFont("Archivo"), "Arial", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
).set(
    body_background_fill="#ECECEC",
    body_background_fill_dark="#ECECEC",
    body_text_color="#000000",
    body_text_color_dark="#000000",
    body_text_color_subdued="#6B6B6B",
    body_text_color_subdued_dark="#6B6B6B",
    background_fill_primary="#FFFFFF",
    background_fill_primary_dark="#FFFFFF",
    background_fill_secondary="#F4F4F4",
    background_fill_secondary_dark="#F4F4F4",
    block_background_fill="#FFFFFF",
    block_background_fill_dark="#FFFFFF",
    block_border_color="#D8D8D8",
    block_border_color_dark="#D8D8D8",
    block_border_width="1px",
    block_label_text_color="#000000",
    block_label_text_color_dark="#000000",
    block_label_background_fill="#FFFFFF",
    block_label_background_fill_dark="#FFFFFF",
    block_title_text_color="#000000",
    block_title_text_color_dark="#000000",
    block_radius="0px",
    block_shadow="none",
    block_shadow_dark="none",
    input_background_fill="#FFFFFF",
    input_background_fill_dark="#FFFFFF",
    input_border_color="#000000",
    input_border_color_dark="#000000",
    input_border_color_focus="#5F2EEA",
    input_border_color_focus_dark="#5F2EEA",
    input_shadow="none",
    input_shadow_dark="none",
    input_shadow_focus="none",
    input_shadow_focus_dark="none",
    shadow_drop="none",
    shadow_drop_lg="none",
    button_primary_background_fill="#5F2EEA",
    button_primary_background_fill_dark="#5F2EEA",
    button_primary_background_fill_hover="#000000",
    button_primary_background_fill_hover_dark="#000000",
    button_primary_text_color="#FFFFFF",
    button_primary_text_color_dark="#FFFFFF",
    button_primary_border_color="#5F2EEA",
    button_primary_border_color_dark="#5F2EEA",
    button_secondary_background_fill="#FFFFFF",
    button_secondary_background_fill_dark="#FFFFFF",
    button_secondary_text_color="#000000",
    button_secondary_text_color_dark="#000000",
    button_secondary_border_color="#000000",
    button_secondary_border_color_dark="#000000",
)

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Archivo:wght@400;500;700;800&family=JetBrains+Mono:wght@400;500&display=swap');

:root {
    --ink: #000000;
    --paper: #FFFFFF;
    --bg: #ECECEC;
    --accent: #5F2EEA;
    --grey: #6B6B6B;
    --hairline: #D8D8D8;
    --grotesk: 'Helvetica Neue', Helvetica, 'Archivo', Arial, sans-serif;
    --mono: 'JetBrains Mono', ui-monospace, monospace;
}

.gradio-container {
    max-width: 100% !important;
    width: 100% !important;
    margin: 0 !important;
    padding-left: clamp(20px, 5vw, 72px) !important;
    padding-right: clamp(20px, 5vw, 72px) !important;
    background: var(--bg) !important;
}

/* Kill every radius and shadow — flat ink on paper */
.gradio-container *,
.gradio-container *::before,
.gradio-container *::after {
    border-radius: 0 !important;
    box-shadow: none !important;
    text-shadow: none !important;
}

/* =========================================================================
   MASTHEAD — wordmark left, tabular meta right, heavy rule below
   ========================================================================= */
#hero { padding: 22px 0 0; margin-bottom: 18px; }
#hero .masthead {
    display: flex;
    justify-content: space-between;
    align-items: flex-end;
    gap: 32px;
    flex-wrap: wrap;
    padding-bottom: 8px;
}
#hero h1 {
    font-family: __TITLE_FAMILY__;
    font-weight: 700;
    font-size: clamp(44px, 8.5vw, 108px);
    line-height: 0.82;
    letter-spacing: -0.05em;
    margin: 0;
    color: var(--ink);
}
#hero h1 .stop { color: var(--accent); }
#hero .meta {
    font-family: var(--grotesk);
    font-size: 11px;
    font-weight: 500;
    text-transform: uppercase;
    letter-spacing: 0.14em;
    line-height: 1;
    color: var(--ink);
    min-width: 280px;
    flex: 0 1 340px;
}
#hero .meta .row {
    display: flex;
    justify-content: space-between;
    gap: 24px;
    padding: 7px 0;
    border-top: 1px solid var(--ink);
}
#hero .meta .row:first-child { border-top: none; }
#hero .meta .k { color: var(--grey); }
#hero .meta .v { text-align: right; }

/* Component labels */
.gradio-container label span,
.gradio-container .gr-check-radio label span {
    font-family: var(--grotesk) !important;
    text-transform: uppercase;
    letter-spacing: 0.12em;
    font-size: 10px !important;
    font-weight: 500 !important;
    color: var(--ink) !important;
}

/* Inputs: ink rules on paper */
.gradio-container textarea,
.gradio-container input[type='text'],
.gradio-container input[type='number'] {
    font-family: var(--grotesk) !important;
    font-size: 14px !important;
    border: 1px solid var(--ink) !important;
    background: var(--paper) !important;
}
.gradio-container textarea:focus,
.gradio-container input:focus {
    border-color: var(--accent) !important;
    outline: 2px solid var(--accent) !important;
    outline-offset: -1px;
}

/* Primary CTA — violet plate, uppercase, no ornament */
.go-btn {
    text-transform: uppercase !important;
    letter-spacing: 0.22em !important;
    font-family: var(--grotesk) !important;
    font-weight: 700 !important;
    font-size: 12px !important;
    min-height: 48px !important;
    border: none !important;
    transition: background-color .12s ease !important;
}
.go-btn:hover { background: var(--ink) !important; }
.go-btn:focus-visible { outline: 2px solid var(--ink) !important; outline-offset: 2px; }

/* Compaction — single-screen fit */
#hero ~ .row, .gradio-container .gap { gap: 12px !important; }
.gradio-container .padded, .gradio-container .block { padding: 10px !important; }
.gradio-container .block { border-color: var(--hairline) !important; }

/* Control image dropzone — icon only, no instruction text */
#control-image .wrap {
    font-size: 0 !important;
    line-height: 0 !important;
    gap: 0 !important;
}
#control-image .wrap .or { display: none !important; }
#control-image .wrap svg { width: 28px !important; height: 28px !important; }

/* Compact output blocks (download + preview) */
#otf-file { max-height: 84px !important; overflow: hidden; }
#otf-file .unpadded_box, #otf-file .empty { min-height: 56px !important; height: 56px !important; }
#preview-img .unpadded_box, #preview-img .empty { min-height: 96px !important; }

/* Progress — single violet gauge inside the grids block */
.gradio-container .progress-bar { background: var(--accent) !important; }
#grids { min-height: 240px; }
#grids .progress-text, #grids .progress-level-inner, #grids .meta-text, #grids .meta-text-center {
    font-family: var(--mono) !important;
    font-size: 12px !important;
    color: var(--ink) !important;
}

@media (prefers-reduced-motion: reduce) {
    .gradio-container * { transition: none !important; animation: none !important; }
}

footer { display: none !important; }
"""

# Inject the chosen title font-family + (optional) embedded @font-face.
CSS = CSS.replace("__TITLE_FAMILY__", TITLE_FAMILY) + TITLE_FONT_FACE

with gr.Blocks(title="Typotopia", theme=THEME, css=CSS) as demo:
    gr.HTML(
        """
        <div id="hero">
          <div class="masthead">
            <h1>Typotopia<span class="stop">.</span></h1>
            <div class="meta">
              <div class="row"><span class="k">Input</span><span class="v">Prompt / Control image</span></div>
              <div class="row"><span class="k">Model</span><span class="v">FLUX.2-klein + LoRA</span></div>
              <div class="row"><span class="k">Output</span><span class="v">OpenType · GPOS kerning</span></div>
            </div>
          </div>
        </div>
        """
    )
    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="e.g. a bold geometric sans inspired by brutalist concrete signage",
                lines=2,
            )
            control_image = gr.Image(label="Control image (optional)", type="pil", height=140, elem_id="control-image")
            btn = gr.Button("Generate font", variant="primary", elem_classes=["go-btn"])
        with gr.Column(scale=1):
            otf_file = gr.File(label="Download .otf", height=56, elem_id="otf-file")
            preview = gr.Image(label="Preview", type="pil", height=150, elem_id="preview-img")
            gallery = gr.Gallery(label="Generated grids", columns=3, height=150, object_fit="contain", elem_id="grids")

    btn.click(
        generate,
        inputs=[prompt, control_image],
        outputs=[gallery, preview, otf_file],
        show_progress="full",
        show_progress_on=gallery,
    )

if __name__ == "__main__":
    demo.queue().launch()