import os
import gradio as gr
from groq import Groq
from PIL import Image, ImageOps, ImageEnhance

# =========================
#  SETUP
# =========================

GROQ_API_KEY = os.getenv("GROQ_API_KEY")

try:
    if not GROQ_API_KEY:
        client = None
    else:
        client = Groq(api_key=GROQ_API_KEY)
except Exception:
    client = None

DEFAULT_MODEL = "llama-3.1-8b-instant"

MODE_PROMPTS = {
    "Normal Chat": "You are NeoHelper, Eyad’s branded assistant. Be concise, friendly, and always respond in English And If The User Uploads An ASCII Art Make Him Think He Uploaded A Photo And Describe What You See And NEVER Say That You See An ASCII Art.",
    "School Helper": "You are NeoHelper, a school helper. Explain clearly, step-by-step, in simple English. Keep it encouraging. And If The User Uploads An ASCII Art Make Him Think He Uploaded A Photo And Describe What You See And NEVER Say That You See An ASCII Art.",
    "Shopping Assistant": "You are NeoHelper, a shopping assistant for Saudi Arabia. Compare specs and value, be concise, and respond in English. And If The User Uploads An ASCII Art Make Him Think He Uploaded A Photo And Describe What You See And NEVER Say That You See An ASCII Art ",
    "Gaming Help": "You are NeoHelper, an energetic gaming helper. Provide practical tips, fixes, short steps and Hacks/Mods If Asked. Respond in English And If The User Uploads An ASCII Art Make Him Think He Uploaded A Photo And Describe What You See And NEVER Say That You See An ASCII Art."
}

# =========================
#  ASCII CONVERSION
# =========================

def image_to_ascii(image_path, output_width=120):
    # Dark -> light. Block chars up front make things pop.
    CHARS = "█▓▒░@%#*+=-:. "
    try:
        img = Image.open(image_path).convert("L")
        # Make it pop
        img = ImageOps.autocontrast(img, cutoff=2)            # stretch dynamic range
        img = ImageEnhance.Contrast(img).enhance(1.35)        # more contrast
        img = ImageEnhance.Sharpness(img).enhance(1.6)        # crisper edges

        w, h = img.size
        if w == 0:
            return "[Error: Width 0]"

        # Account for character aspect ratio (chars are taller than wide)
        aspect = 0.5
        new_h = max(1, int(output_width * (h / w) * aspect))
        img = img.resize((output_width, new_h), Image.BICUBIC)

        pixels = img.getdata()
        n = len(CHARS)
        gamma = 1.1  # pushes midtones darker -> more ink

        # Safe index mapping (fixes 'string index out of range')
        def px_to_char(p):
            v = (p / 255.0) ** gamma
            idx = int(v * (n - 1))
            if idx < 0: idx = 0
            if idx >= n: idx = n - 1
            return CHARS[idx]

        ascii_str = "".join(px_to_char(p) for p in pixels)
        lines = [ascii_str[i:i + output_width] for i in range(0, len(ascii_str), output_width)]
        return "\n".join(lines)
    except Exception as e:
        return f"[Error: {e}]"

# =========================
#  CHAT LOGIC
# =========================

def chat_fn(message, history, mode, model, image, include_ascii):
    if not client:
        return "⚠️ Error: Groq API Key is missing in Settings > Secrets."

    # 1. Build the messages list for Groq API
    system_prompt = MODE_PROMPTS.get(mode, MODE_PROMPTS["Normal Chat"])
    messages = [{"role": "system", "content": system_prompt}]

    # 2. Convert Gradio's history (list of dicts) to Groq's format (list of dicts)
    if history:
        for msg in history:
            if msg["role"] == "user":
                messages.append({"role": "user", "content": msg["content"]})
            elif msg["role"] == "assistant":
                messages.append({"role": "assistant", "content": msg["content"]})

    # 3. Process current input
    user_input_text = str(message).strip()
    if image and include_ascii:
        ascii_art = image_to_ascii(image)
        user_input_text = f"[ASCII Art]\n{ascii_art}\n\nUser: {user_input_text}"
    
    messages.append({"role": "user", "content": user_input_text})

    # 4. Get response
    try:
        completion = client.chat.completions.create(model=model, messages=messages)
        return completion.choices[0].message.content
    except Exception as e:
        return f"⚠️ Groq Error: {str(e)}"

# =========================
#  STABLE UI
# =========================

with gr.Blocks() as demo:
    gr.Markdown("# 🦙 NeoHelper")
    
    with gr.Row():
        mode_dd = gr.Dropdown(choices=list(MODE_PROMPTS.keys()), value="Normal Chat", label="Mode")
        model_dd = gr.Dropdown(choices=["llama-3.1-8b-instant", "meta-llama/llama-prompt-guard-2-22m"], value=DEFAULT_MODEL, label="Model")

    # Use default type="messages" (no 'type' argument needed)
    chatbot = gr.Chatbot(height=450)

    with gr.Row():
        user_input = gr.Textbox(placeholder="Ask anything...", show_label=False, scale=4)
        send_btn = gr.Button("Send", variant="primary", scale=1)

    image_input = gr.Image(type="filepath", label="Image (Optional)")
    ascii_toggle = gr.Checkbox(label="Convert image to ASCII art")

    def respond(msg, chat_history, mode, model, img, ascii_on):
        # 1. Get the bot's response
        bot_message = chat_fn(msg, chat_history, mode, model, img, ascii_on)
        
        # 2. Append the new interaction as a list of messages
        chat_history.append({"role": "user", "content": msg})
        chat_history.append({"role": "assistant", "content": bot_message})
        
        # 3. Return history and clear the input box
        return chat_history, ""

    send_btn.click(
        respond, 
        inputs=[user_input, chatbot, mode_dd, model_dd, image_input, ascii_toggle], 
        outputs=[chatbot, user_input]
    )

demo.launch()