import os
import gradio as gr
from groq import Groq
from dotenv import load_dotenv
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch

# ===============================
# Load API Key from .env
# ===============================
load_dotenv()
GROQ_API_KEY = os.getenv("Freelance_key")

client = Groq(api_key=GROQ_API_KEY)

# ===============================
# AI Response Function
# ===============================
def generate_response(user_input, chat_history):
    if chat_history is None:
        chat_history = []

    # Prepare messages for Groq AI
    messages = [{"role": "system", "content": "You are a professional AI assistant."}]
    for msg in chat_history:
        messages.append(msg)
    messages.append({"role": "user", "content": user_input})

    # Call Groq AI
    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages,
        temperature=0.7
    )
    answer = response.choices[0].message.content

    # Append messages in new format
    chat_history.append({"role": "user", "content": user_input})
    chat_history.append({"role": "assistant", "content": answer})

    return chat_history, chat_history

# ===============================
# PDF Download Function
# ===============================
def download_pdf(chat_history):
    file_path = "AI_Agent_Chat.pdf"
    doc = SimpleDocTemplate(file_path)
    elements = []
    styles = getSampleStyleSheet()

    for message in chat_history:
        role = message["role"]
        content = message["content"]
        elements.append(Paragraph(f"<b>{role.capitalize()}:</b> {content}", styles["Normal"]))
        elements.append(Spacer(1, 0.3 * inch))

    doc.build(elements)
    return file_path

# ===============================
# Gradio UI
# ===============================
with gr.Blocks() as demo:
    # Instagram style header with emoji
    gr.Markdown("💜 **Instagram Style AI Agent**")

    chatbot = gr.Chatbot()
    state = gr.State([])

    user_input = gr.Textbox(label="Type your message here...")
    generate_btn = gr.Button("Send")
    download_btn = gr.Button("Download Chat as PDF")

    generate_btn.click(
        generate_response,
        inputs=[user_input, state],
        outputs=[chatbot, state]
    )

    download_btn.click(
        download_pdf,
        inputs=[state],
        outputs=gr.File()
    )

demo.launch()