import base64, io, os, json
import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image

# --- HF Inference Client -----------------------------------------------------
client = InferenceClient(
    provider="nebius",
    api_key=os.environ.get("HF_TOKEN", ""),
)

SYSTEM_PROMPT = (
    "Return ONLY a minified JSON object with keys: "
    "name, title, email, link — in that order."
)

# --- Utility helpers ---------------------------------------------------------

def pil_to_data_url(pil_img: Image.Image) -> str:
    """Convert PIL image to a base‑64 data‑URL readable by the vision model."""
    buf = io.BytesIO()
    pil_img.convert("RGB").save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode()
    return f"data:image/png;base64,{b64}"


def make_html_card(json_str: str) -> str:
    """Generate a colourful HTML business‑card preview from JSON."""
    try:
        data = json.loads(json_str)
    except json.JSONDecodeError:
        return "<p style='color:red'>🤔 Hmm, the model did not return valid JSON.</p>"

    name  = data.get("name", "")
    title = data.get("title", "")
    email = data.get("email", "")
    link  = data.get("link", "")

    return f"""
    <div style='font-family:Inter,Arial,sans-serif;border-radius:16px;padding:24px;width:280px;
                background:linear-gradient(135deg,#8BC6EC 0%,#9599E2 100%);color:#fff;
                box-shadow:0 4px 14px rgba(0,0,0,.2)'>
        <h2 style='margin:0 0 6px 0;font-size:24px;'>{name}</h2>
        <h4 style='margin:0 0 12px 0;font-weight:400;'>{title}</h4>
        <p style='margin:0 0 8px 0;font-size:14px;'>📧 {email}</p>
        <p style='margin:0;font-size:14px;'>🔗 <a style='color:#fff;text-decoration:underline;' href='{link}' target='_blank'>{link}</a></p>
    </div>
    """


# --- Core generation fn ------------------------------------------------------

def generate_business_card(pil_img: Image.Image | None):
    """Return both raw JSON & a pretty HTML preview."""
    if pil_img is None:
        return "// Upload a selfie first 🙂", ""

    data_url = pil_to_data_url(pil_img)

    resp = client.chat.completions.create(
        model="google/gemma-3-27b-it",
        max_tokens=200,
        temperature=0.9,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Generate my JSON business card:"},
                    {"type": "image_url", "image_url": {"url": data_url}},
                ],
            },
        ],
    )

    json_str  = resp.choices[0].message.content.strip().strip("`")
    card_html = make_html_card(json_str)

    return json_str, card_html


# --- Theme -------------------------------------------------------------------

theme = (
    gr.themes.Default(primary_hue="indigo", secondary_hue="purple")
    .set(body_background_fill="linear-gradient(120deg,#f6d365 0%,#fda085 100%)")
)


# --- UI ----------------------------------------------------------------------

with gr.Blocks(title="Business‑Card Agent", theme=theme,
               css=".gradio-container {font-family:'Inter',sans-serif;}") as demo:
    gr.Markdown(
        """
        <h1 style='display:flex;align-items:center;gap:8px;'>📇 Business‑Card Agent</h1>
        <p>Upload a selfie & get a <code>business_card.json</code> draft — plus a live preview!</p>
        """
    )

    with gr.Row():
        selfie = gr.Image(label="Selfie", type="pil", show_label=True)
        with gr.Column():
            code = gr.Code(
                label="business_card.json",
                language="json",
                interactive=False,
                value='{"name":"","title":"","email":"","link":""}'
            )
            card_html = gr.HTML(label="Card Preview", value="")

    gr.Button("✨ Generate", variant="primary").click(
        generate_business_card, selfie, [code, card_html]
    )

    gr.Markdown(
        "<small>Powered by 🤗 <a href='https://huggingface.co/'>HuggingFace</a> — Source on <a href='https://github.com/'>GitHub</a></small>"
    )

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