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
load_dotenv()
GROQ_API_KEY = os.getenv("Freelancing_key")

client = Groq(api_key=GROQ_API_KEY)

# ===============================
# AI RESPONSE FUNCTION (NEW FORMAT)
# ===============================

def generate_response(task, user_input, tone, history):

    if history is None:
        history = []

    prompt = f"""
Task: {task}
Tone: {tone}
User Request: {user_input}
"""

    messages = [
        {"role": "system", "content": "You are a professional AI business assistant."}
    ]

    # Add previous chat correctly
    for message in history:
        messages.append(message)

    messages.append({"role": "user", "content": prompt})

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages,
        temperature=0.7
    )

    answer = response.choices[0].message.content

    # Append in NEW format
    history.append({"role": "user", "content": user_input})
    history.append({"role": "assistant", "content": answer})

    return history, history


# ===============================
# PDF DOWNLOAD FUNCTION
# ===============================

def download_pdf(history):

    file_path = "AI_Agent_Chat.pdf"
    doc = SimpleDocTemplate(file_path)
    elements = []
    styles = getSampleStyleSheet()

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

    doc.build(elements)
    return file_path


# ===============================
# GRADIO UI (NEW FORMAT)
# ===============================

with gr.Blocks() as demo:

    gr.Markdown("# 💰 Freelancing AI Agent Pro")

    chatbot = gr.Chatbot()   # NEW version auto-detects message format
    state = gr.State([])

    task = gr.Dropdown(
        ["Proposal Writer", "Email Writer", "Marketing Post", "Business Idea Generator"],
        label="Select Task"
    )

    tone = gr.Dropdown(
        ["Professional", "Friendly", "Confident", "Persuasive"],
        label="Select Tone"
    )

    user_input = gr.Textbox(label="Enter Your Request")

    generate_btn = gr.Button("Generate")
    download_btn = gr.Button("Download as PDF")

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

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

demo.launch()