import os
import base64
from groq import Groq
import gradio as gr

# -- Groq Client Setup for Text Chats
def get_response(message, history):
    client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

    # System prompt
    messages = [
        {
            "role": "system",
            "content": (
                "You are a UX Ops assistant. Help the design team improve processes. You can:\n"
                "- Collect feedback by topic\n"
                "- Summarize research transcripts\n"
                "- Assist with dev handoff checklists\n"
                "- Ask follow-up questions when needed"
            )
        }
    ]

    # Clean Gradio history
    clean_history = [
        {"role": turn["role"], "content": turn["content"]}
        for turn in history
        if isinstance(turn, dict) and "role" in turn and "content" in turn
    ]
    messages.extend(clean_history)

    # Append new user message
    messages.append({"role": "user", "content": message})

    # Call the API
    chat_completion = client.chat.completions.create(
        messages=messages,
        model="meta-llama/llama-4-scout-17b-16e-instruct"
    )
    return chat_completion.choices[0].message.content

# -- Other Handlers
def summarize_file(file):
    with open(file.name, "r") as f:
        content = f.read()
    return get_response(f"Summarize this user interview:\n{content}", [])

def validate_handoff(checked_items):
    msg = "Here’s the developer handoff readiness:\n"
    status = []
    status.append("✅ Components used from design system" if "components" in checked_items else "⚠️ Components not confirmed")
    status.append("✅ Developer notes added" if "notes" in checked_items else "⚠️ Notes missing")
    status.append("✅ Dev links and specs included" if "links" in checked_items else "⚠️ Links to specs missing")
    return "\n".join([msg] + status)

def structure_feedback(layout, copy, interaction):
    return get_response(
        f"Collect feedback:\n- Layout: {layout}\n- Copy: {copy}\n- Interaction: {interaction}",
        []
    )

# -- Design Image Feedback with Multimodal Support
def provide_design_feedback(image_path):
    # Encode image to Data URL
    with open(image_path, "rb") as img_f:
        encoded = base64.b64encode(img_f.read()).decode("utf-8")
    ext = os.path.splitext(image_path)[1].lstrip(".").lower()
    data_url = f"data:image/{ext};base64,{encoded}"

    # Multimodal user message
    user_content = [
        {"type": "text", "text": (
            "Audit the UX of this design—provide feedback on hierarchy, color, typography, "
            "accessibility, and usability."
        )},
        {"type": "image_url", "image_url": {"url": data_url}}
    ]

    # Build messages
    client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
    messages = [
        {"role": "system", "content": "You are a UX Ops assistant. Give actionable, concise design feedback."},
        {"role": "user", "content": user_content}
    ]

    # Call Groq multimodal completion
    chat_completion = client.chat.completions.create(
        model="meta-llama/llama-4-maverick-17b-128e-instruct",
        messages=messages,
        temperature=1.0,
        max_completion_tokens=1024,
        top_p=1.0
    )
    return chat_completion.choices[0].message.content

# -- Main App UI
with gr.Blocks(title="UX Ops Assistant") as demo:
    gr.Markdown("## 🎯 UX Design Operations Dashboard", elem_classes="centered")

    # Chat Tab
    with gr.Tab("🧠 Ask Assistant"):
        chatbot = gr.Chatbot(type="messages", height=400)
        user_input = gr.Textbox(placeholder="Ask about UX process, feedback, etc.")

        def submit_chat(msg, history):
            reply = get_response(msg, history)
            new_history = history + [
                {"role": "user", "content": msg},
                {"role": "assistant", "content": reply}
            ]
            return "", new_history

        user_input.submit(
            fn=submit_chat,
            inputs=[user_input, chatbot],
            outputs=[user_input, chatbot]
        )

    # Research Upload Tab
    with gr.Tab("📂 Upload Research"):
        file_input = gr.File(label="Upload transcript (.txt)")
        summary_output = gr.Textbox(label="Summary", lines=6)
        file_input.change(fn=summarize_file, inputs=file_input, outputs=summary_output)

    # Dev Handoff Checklist Tab
    with gr.Tab("✅ Dev Handoff Checklist"):
        checklist = gr.CheckboxGroup(
            choices=["components", "notes", "links"],
            label="What's included in this design?"
        )
        validate_btn = gr.Button("Validate Handoff Readiness")
        result = gr.Textbox(label="Result", lines=4)
        validate_btn.click(fn=validate_handoff, inputs=checklist, outputs=result)

    # Feedback Organization Tab
    with gr.Tab("✏️ Collect Design Feedback"):
        layout_fb = gr.Textbox(label="Layout")
        copy_fb = gr.Textbox(label="Copy/Wording")
        interaction_fb = gr.Textbox(label="Interactions")
        fb_btn = gr.Button("Organize Feedback")
        fb_output = gr.Textbox(label="Organized Feedback", lines=6)
        fb_btn.click(
            fn=structure_feedback,
            inputs=[layout_fb, copy_fb, interaction_fb],
            outputs=fb_output
        )

    # Design Image Feedback Tab
    with gr.Tab("🖼️ Upload Design Feedback"):
        design_image = gr.Image(type="filepath", label="Upload Design (JPG, PNG)")
        image_preview = gr.Image(interactive=False, label="Preview")
        feedback_btn2 = gr.Button("Get Design Feedback")
        design_feedback_output = gr.Textbox(label="Design Feedback", lines=10)

        def display_and_feedback(image_path):
            preview = image_path
            feedback = provide_design_feedback(image_path)
            return preview, feedback

        feedback_btn2.click(
            fn=display_and_feedback,
            inputs=design_image,
            outputs=[image_preview, design_feedback_output]
        )

    demo.launch()