import os
import json
import re
import io
import base64
import streamlit as st
from PIL import Image
from groq import Groq


# --- Helper Functions ---
def text_to_json(user_prompt, groq_client, model_name):
    """Convert natural language to structured JSON."""
    prompt = f"""
    You are a helpful assistant that converts ideas into structured JSON.
    The JSON must be valid and minimal.

    Input: "{user_prompt}"

    Return ONLY valid JSON (no explanations or code comments).
    """
    try:
        response = groq_client.chat.completions.create(
            model=model_name,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3,
            max_tokens=800,
        )
        raw_output = response.choices[0].message.content.strip()

        try:
            parsed = json.loads(raw_output)
        except json.JSONDecodeError:
            match = re.search(r"\{[\s\S]*\}", raw_output)
            if match:
                parsed = json.loads(match.group(0))
            else:
                raise
        return json.dumps(parsed, indent=2)
    except Exception as e:
        return f"Error: {str(e)}"


def json_to_text(json_text, groq_client, model_name):
    """Convert JSON schema or data back into descriptive text."""
    prompt = f"""
    You are an assistant that explains JSON data in clear, natural language.
    Turn the following JSON into a human-readable description or creative prompt.

    JSON Input:
    {json_text}

    Return only the descriptive text.
    """
    try:
        response = groq_client.chat.completions.create(
            model=model_name,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.4,
            max_tokens=600,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"Error: {str(e)}"


def image_to_prompt(image, groq_client, model_name):
    """Generate a text prompt from an uploaded image using text-only fallback."""
    try:
        img_buffer = io.BytesIO()
        image.save(img_buffer, format="PNG")
        img_b64 = base64.b64encode(img_buffer.getvalue()).decode()

        prompt = (
            "You are an assistant that receives base64-encoded image data and describes it in vivid detail. "
            "Analyze what the image likely depicts and return a creative, natural-language prompt. "
            "Here is the base64 image data (truncated for safety):\n"
            f"{img_b64[:4000]}..."
        )

        response = groq_client.chat.completions.create(
            model=model_name,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.5,
            max_tokens=500,
        )
        return response.choices[0].message.content.strip()

    except Exception as e:
        return f"Error: {str(e)}"


# --- Streamlit App ---
def main():
    st.set_page_config(page_title="Prompt ↔ JSON ↔ Image Generator", page_icon="🧠", layout="wide")

    st.markdown("""
    <style>
        .block-container { padding-left: 2rem; padding-right: 2rem; max-width: 1200px; }
        textarea { font-family: monospace; }
    </style>
    """, unsafe_allow_html=True)

    st.title("🧠 Prompt ↔ JSON ↔ Image Generator")
    st.markdown("""
    Convert freely between **text**, **JSON**, and **image-based prompts** —  
    perfect for AI development, creative workflows, and generative pipelines.
    """)

    # Sidebar
    st.sidebar.image("/app/src/2.png")
    st.sidebar.image("/app/src/1.PNG", caption="Powered by Groq")
    st.sidebar.markdown("Build by **DW** ⚡")

    groq_api_key = st.secrets.get("GROQ_API_KEY", "") or os.getenv("GROQ_API_KEY")
    if not groq_api_key:
        st.error("Missing GROQ_API_KEY in environment or Streamlit secrets.")
        st.stop()

    groq_client = Groq(api_key=groq_api_key)

    model_name = st.sidebar.selectbox(
        "Select Model:",
        ["openai/gpt-oss-120b","llama-3.3-70b-versatile", "mixtral-8x7b-32768"]
    )

    mode = st.sidebar.radio(
        "Choose Direction:",
        ["Text → JSON", "JSON → Text", "🖼️ Image → Prompt"]
    )

    # --- Mode 1: Text → JSON ---
    if mode == "Text → JSON":
        st.subheader("📝 Describe your idea or task")
        user_prompt = st.text_area(
            "Enter natural language description:",
            height=180,
            placeholder="Example: Create a JSON schema for a quest system with title, reward, and difficulty."
        )

        if st.button("Generate JSON"):
            if not user_prompt.strip():
                st.warning("Please enter a prompt first.")
                return

            with st.spinner("Generating JSON..."):
                result = text_to_json(user_prompt, groq_client, model_name)

            if result.startswith("Error"):
                st.error(result)
            else:
                st.subheader("🧩 Generated JSON:")
                st.code(result, language="json")
                st.download_button(
                    "💾 Download JSON",
                    data=result,
                    file_name="generated_prompt.json",
                    mime="application/json"
                )

    # --- Mode 2: JSON → Text ---
    elif mode == "JSON → Text":
        st.subheader("🧩 Paste your JSON input")
        json_text = st.text_area(
            "Enter JSON structure:",
            height=250,
            placeholder='{\n  "task": "daily_affirmation",\n  "fields": {"message": "string", "mood": "integer"}\n}'
        )

        if st.button("Generate Description"):
            if not json_text.strip():
                st.warning("Please enter JSON first.")
                return

            with st.spinner("Converting JSON to prompt..."):
                result = json_to_text(json_text, groq_client, model_name)

            if result.startswith("Error"):
                st.error(result)
            else:
                st.subheader("📝 Generated Description:")
                st.write(result)
                st.download_button(
                    "💾 Download Text",
                    data=result,
                    file_name="generated_prompt.txt",
                    mime="text/plain"
                )

    # --- Mode 3: Image → Prompt ---
    elif mode == "🖼️ Image → Prompt":
        st.subheader("🖼️ Upload an image to convert to a text prompt")
        uploaded_image = st.file_uploader("Upload an image (PNG, JPG)", type=["png", "jpg", "jpeg"])

        if uploaded_image:
            image = Image.open(uploaded_image)
            st.image(image, caption="Uploaded Image")

            if st.button("Generate Prompt from Image"):
                with st.spinner("Analyzing image..."):
                    result = image_to_prompt(image, groq_client, model_name)

                if result.startswith("Error"):
                    st.error(result)
                else:
                    st.subheader("🧠 Generated Prompt:")
                    st.text_area("Prompt Text", result, height=200)
                    st.download_button(
                        "💾 Download Prompt",
                        data=result,
                        file_name="image_prompt.txt",
                        mime="text/plain"
                    )


if __name__ == "__main__":
    main()