import os
import json
import asyncio
from datetime import datetime
import gradio as gr
import httpx

from websearch import PerplexityClient, parse_perplexity_response

# load .env file
from dotenv import load_dotenv
load_dotenv()

# Initialize Perplexity client
perplexity_client = PerplexityClient(
    api_key=os.environ["PERPLEXITY_AUTH_TOKEN"]
    )

async def query_api(query: str, website: str, after_date: str) -> tuple:
    if website:
        query = query + f" \"site:{website}\""
    if after_date:
        query = query + f" \"after:{after_date}\""
    print("query: ", query)
    try:
        # Fetch response from Perplexity
        response = await perplexity_client.generate_response(query)
        
        # Parse the response
        parsed_response = parse_perplexity_response(response)
        
        response_data = json.dumps(parsed_response)
                
        # Parse the accumulated response data as JSON
        response_json = json.loads(response_data)
        
        # Extract content and citations from the response JSON
        content = response_json.get("content", "")
        citations = response_json.get("citations", [])
        
        # Beautify content using Markdown formatting
        beautified_content = f"# Search Results\n\n{content}"
        
        # Beautify citations by adding Markdown links
        beautified_citations = "# Citations/Sources\n\n"
        for i, citation in enumerate(citations, start=1):
            beautified_citations += f"{i}. [{citation}]({citation})\n"
        
        # Yield the beautified content and citations
        yield beautified_content, beautified_citations
    except httpx.TimeoutException:
        yield "# Request Timeout\n\nRequest timed out. Please try again later.", ""
    except httpx.HTTPStatusError as e:
        yield f"# HTTP Error\n\nHTTP error occurred: {e}", ""
    except Exception as e:
        yield f"# Error\n\nAn error occurred: {e}", ""

# Create Gradio interface
with gr.Blocks(css=".gradio-container { background-color: #f5f5f5; padding: 20px; border-radius: 10px; }", theme=gr.themes.Citrus()) as demo:
    gr.Markdown("# Web Search Application")
    
    with gr.Row():
        with gr.Column(
            render=True,
            show_progress=True
        ):
            query = gr.Textbox(
                label="Enter your query",
                placeholder="Type your search query here...",
                lines=2,
                max_lines=20,
                value="",
                elem_id="query-input"
            )
            website = gr.Textbox(
                label="Single website search",
                placeholder="website url",
                lines=1,
                max_lines=20,
                value="",
                elem_id="query-input"
            )
            after_date = gr.Textbox(
                label="After date",
                placeholder="date format: YYYY-MM-DD",
                lines=1,
                max_lines=2,
                value="",
                elem_id="query-input"
            )
            submit_button = gr.Button("Search")
        
        with gr.Column(
            render=True,
            show_progress=True
        ):
            output_content = gr.Markdown(
                label="Response Content",
                value="",
                elem_id="response-content",
                height="600px",
                visible=True,
                show_label=True
                
            )
            output_citations = gr.Markdown(
                label="Citations",
                value="",
                elem_id="response-citations",
                height="200px",
                visible=True,
                show_label=True
            )
    
    # Set up event listener
    submit_button.click(query_api, inputs=[query, website, after_date], outputs=[output_content, output_citations])
    
    gr.Markdown("Powered by FastAPI and Gradio")

# Launch the Gradio application
demo.launch()