import base64
import gradio as gr
from huggingface_hub import InferenceClient

# Free serverless models
TEXT_MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"   # strong at PHP / WordPress code
VISION_MODEL = "Qwen/Qwen2.5-VL-72B-Instruct"     # can read screenshots

client = InferenceClient()

SYSTEM_PROMPT = (
    "You are a senior WordPress plugin developer and web developer assistant. "
    "You help write, debug, and explain PHP, WordPress plugin code (hooks, filters, "
    "shortcodes, admin settings pages, Media Library integration), HTML, CSS, and "
    "JavaScript. When shown a screenshot of a website, describe the layout/section "
    "the user points to, then rebuild it as real HTML/CSS/PHP as closely as possible. "
    "Give complete, working code. When code is long, deliver it in full, installable "
    "form (ready to save as a .php file)."
)

def image_to_data_url(img_path):
    with open(img_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("utf-8")
    ext = img_path.split(".")[-1].lower()
    mime = "jpeg" if ext in ("jpg", "jpeg") else ext
    return f"data:image/{mime};base64,{b64}"

def chat(message, history):
    # message can be a dict with "text" and "files" when multimodal input is used
    text = message.get("text", "") if isinstance(message, dict) else message
    files = message.get("files", []) if isinstance(message, dict) else []

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in history:
        if isinstance(turn, dict):
            role = turn.get("role")
            content = turn.get("content")
            if role and content:
                messages.append({"role": role, "content": content})

    if files:
        # Vision path: send image + text together to the VL model
        content = [{"type": "text", "text": text or "What is in this screenshot? Rebuild the relevant section as code."}]
        for f in files:
            content.append({"type": "image_url", "image_url": {"url": image_to_data_url(f)}})
        messages.append({"role": "user", "content": content})
        model = VISION_MODEL
    else:
        messages.append({"role": "user", "content": text})
        model = TEXT_MODEL

    try:
        response = client.chat_completion(
            model=model,
            messages=messages,
            max_tokens=2048,
            temperature=0.3,
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"⚠️ Error calling the model: {e}"

demo = gr.ChatInterface(
    fn=chat,
    title="Free WordPress Dev Assistant (with screenshot support)",
    description=(
        "Free coding helper for WordPress plugins, PHP, HTML, and JS. "
        "Attach a screenshot and ask it to rebuild that section as code. "
        "Runs on Qwen2.5-Coder and Qwen2.5-VL via Hugging Face's free Inference API."
    ),
    multimodal=True,
)

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