import os
import io
import random
import base64
import json
import gradio as gr
from huggingface_hub import HfFileSystem
from PIL import Image

DATASET_ID = "roxopastel/999"
HF_TOKEN = os.environ.get("HF_TOKEN", "")

fs = HfFileSystem(token=HF_TOKEN if HF_TOKEN else None)


def list_images() -> list[str]:
    try:
        files = fs.ls(f"datasets/{DATASET_ID}", detail=False)
        return [f for f in files if f.lower().endswith((".png", ".jpg", ".jpeg"))]
    except Exception:
        return []


def read_captions() -> dict[str, str]:
    captions = {}
    try:
        path = f"datasets/{DATASET_ID}/metadata.jsonl"
        with fs.open(path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                fname = obj.get("file_name", "")
                cap = obj.get("caption", "")
                if fname:
                    captions[fname] = cap
    except Exception:
        pass
    return captions


def get_random_image() -> str:
    images = list_images()
    if not images:
        return _error_html("⚠️ Não foi possível acessar o dataset. Verifique o HF_TOKEN.")

    path = random.choice(images)
    file_name = path.split("/")[-1]

    captions = read_captions()
    caption = captions.get(file_name, "")
    if not caption:
        caption = f"#{file_name.split('.')[0]}"

    try:
        with fs.open(path, "rb") as f:
            img_bytes = f.read()

        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=92)
        b64 = base64.b64encode(buf.getvalue()).decode()

        return f"""
        <div style="
            display: flex;
            flex-direction: column;
            align-items: center;
            padding: 0 1rem;
        ">
            <div style="
                background: #ffffff;
                padding: 10px 10px 0 10px;
                box-shadow: 0 2px 8px rgba(0,0,0,0.10), 0 12px 32px rgba(0,0,0,0.09);
                border-radius: 2px;
                width: 100%;
                max-width: 600px;
            ">
                <img
                    src="data:image/jpeg;base64,{b64}"
                    style="
                        display: block;
                        width: 100%;
                        aspect-ratio: 1/1;
                        object-fit: contain;
                    "
                />
                <div style="
                    padding: 14px 12px 18px;
                    text-align: center;
                    font-family: 'Instrument Serif', serif;
                    font-style: italic;
                    font-size: clamp(1.1rem, 3.5vw, 1.4rem);
                    color: #111111;
                    line-height: 1.4;
                ">
                    {caption}
                </div>
            </div>
        </div>
        """

    except Exception as e:
        return _error_html(f"⚠️ Erro ao carregar imagem: {e}")


def _error_html(msg: str) -> str:
    return f'<p style="text-align:center;font-family:monospace;color:#c45c3a;">{msg}</p>'


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Instrument+Serif:ital@0;1&family=Space+Mono:wght@400;700&display=swap');

:root {
    --cream: #faf6f0;
    --ink:   #1a1310;
    --rust:  #c45c3a;
}

body,
.gradio-container,
.gradio-container > .main,
.gradio-container > .main > .wrap,
.contain {
    background: var(--cream) !important;
    padding: 0 !important;
    margin: 0 auto !important;
    max-width: 100% !important;
    gap: 0 !important;
}

.gradio-container > .main > .wrap > .gap,
.block {
    display: flex !important;
    flex-direction: column !important;
    align-items: center !important;
    gap: 0 !important;
}

/* ── title ── */
#header {
    width: 100%;
    text-align: center;
    padding: 1.4rem 1rem 0.5rem;
}

#header h1 {
    font-family: 'Instrument Serif', serif !important;
    font-size: clamp(2.2rem, 8vw, 4rem) !important;
    font-weight: 400 !important;
    color: var(--ink) !important;
    letter-spacing: -0.02em;
    line-height: 1;
    margin: 0 !important;
}

/* ── button ── */
#spin-btn {
    clip-path: polygon(
        10px 0%, calc(100% - 10px) 0%,
        100% 10px, 100% calc(100% - 10px),
        calc(100% - 10px) 100%, 10px 100%,
        0% calc(100% - 10px), 0% 10px
    ) !important;
    width: 52px !important;
    height: 52px !important;
    min-width: 52px !important;
    max-width: 52px !important;
    min-height: 52px !important;
    max-height: 52px !important;
    flex: 0 0 52px !important;
    border-radius: 0 !important;
    background: var(--ink) !important;
    border: none !important;
    font-size: 1.3rem !important;
    padding: 0 !important;
    margin: 0.6rem auto 0.8rem !important;
    cursor: pointer !important;
    transition: background 0.15s, transform 0.1s !important;
    box-shadow: none !important;
    filter: drop-shadow(3px 3px 0px var(--rust)) !important;
}

#spin-btn:hover {
    background: var(--rust) !important;
    transform: scale(1.08) !important;
}

#spin-btn:active {
    transform: scale(0.92) !important;
}

/* ── html output block ── */
#polaroid-out {
    width: 100% !important;
    padding: 0 !important;
    margin: 0 !important;
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
}

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

with gr.Blocks(title="999 Doodles") as demo:

    gr.HTML('<div id="header"><h1>999 Doodles</h1></div>')

    spin_btn = gr.Button("✨", elem_id="spin-btn", variant="primary")

    polaroid = gr.HTML(
        value="",
        elem_id="polaroid-out",
    )

    spin_btn.click(fn=get_random_image, inputs=[], outputs=[polaroid])

demo.launch(css=CSS)