import gradio as gr
from ai_generate import generate_rag

def process_input(topic, length, tone, format_, pdfs):
    # Construct the prompt
    prompt = f"Write a {format_} about {topic} in about {length} words and a {tone} tone."
    print(prompt)
    # Generate the text and citations using RAG
    rag_output = generate_rag(
        prompt=prompt,
        topic=topic,
        model="OpenAI GPT 4o",  # Replace with your model name or path
        url_content=None,
        path=pdfs,
        temperature=1.0,
        max_length=2048,
        api_key="",  # Add your API key if necessary
        sys_message=""
    )
    
    return rag_output


def generate(
    prompt: str,
    topic: str,
    model: str,
    url_content: dict,
    path: list[str],
    temperature: float = 1.0,
    max_length: int = 2048,
    api_key: str = "",
    sys_message="",
):
    return generate_rag(prompt, topic, model, url_content, path, temperature, max_length, api_key, sys_message)

def create_app():
    with gr.Blocks() as app:
        with gr.Row():
            with gr.Column(scale=2):
                topic_input = gr.Textbox(
                            label="Topic",
                            placeholder="Enter the main topic of your article",
                            elem_classes="input-highlight-pink",
                        )
                length_input = gr.Slider(
                            minimum=50,
                            maximum=500,
                            step=50,
                            value=200,
                            label="Article Length",
                            elem_classes="input-highlight-pink",
                        )
                tone_input = gr.Dropdown(
                                choices=[
                                    "Formal",
                                    "Informal",
                                    "Technical",
                                    "Conversational",
                                    "Journalistic",
                                    "Academic",
                                    "Creative",
                                ],
                                value="Formal",
                                label="Writing Style",
                                elem_classes="input-highlight-yellow",
                            )
                format_input = gr.Dropdown(
                                choices=[
                                    "Article",
                                    "Essay",
                                    "Blog post",
                                    "Report",
                                    "Research paper",
                                    "News article",
                                    "White paper",
                                    "Email",
                                    "LinkedIn post",
                                    "X (Twitter) post",
                                    "Instagram Video Content",
                                    "TikTok Video Content",
                                    "Facebook post",
                                ],
                                value="Article",
                                label="Format",
                                elem_classes="input-highlight-turquoise",
                            )
    
                pdf_input = gr.File(label="Upload PDFs", file_types=["pdf"], file_count="multiple")
                generate_button = gr.Button("Generate")
            with gr.Column(scale=3):
                generated_text_output = gr.Textbox(label="Generated Text", lines=10)
                citations_output = gr.HTML(label="Citations")
                
                generate_button.click(
                    fn=process_input,
                    inputs=[topic_input, length_input, tone_input, format_input, pdf_input],
                    outputs=[generated_text_output, citations_output]
                )
        
    return app

# Run the app
app = create_app()
app.launch()