from fastapi import FastAPI, HTTPException, Request, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Any
from helper_functions_api import md_to_html, search_brave, fetch_and_extract_content, limit_tokens, together_response, insert_data
import os
from dotenv import load_dotenv, find_dotenv

# Load environment variables from .env file
#load_dotenv("keys.env")

app = FastAPI()
TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
BRAVE_API_KEY = os.getenv('BRAVE_API_KEY')
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
HELICON_API_KEY = os.getenv("HELICON_API_KEY")
SUPABASE_USER = os.environ['SUPABASE_USER']
SUPABASE_PASSWORD = os.environ['SUPABASE_PASSWORD']

# Groq model names
llm_default_small = "llama3-8b-8192"
llm_default_medium = "llama3-70b-8192"

# Together Model names (fallback)
llm_fallback_small = "meta-llama/Llama-3-8b-chat-hf"
llm_fallback_medium = "meta-llama/Llama-3-70b-chat-hf"

SysPromptJson = "You are now in the role of an expert AI who can extract structured information from user request. Both key and value pairs must be in double quotes. You must respond ONLY with a valid JSON file. Do not add any additional comments."
SysPromptList = "You are now in the role of an expert AI who can extract structured information from user request. All elements must be in double quotes. You must respond ONLY with a valid python List. Do not add any additional comments."
SysPromptDefault = "You are an expert AI, complete the given task. Do not add any additional comments."
SysPromptMd = "You are an expert AI who can create a structured report using information provided in the context from user request.The report should be in markdown format consists of markdown tables structured into subtopics. Do not add any additional comments."

sys_prompts = {
    "offline": {
        "Chat": "You are an expert AI, complete the given task. Do not add any additional comments.",
        "Full Text Report": "You are an expert AI who can create a detailed report from user request. The report should be in markdown format. Do not add any additional comments.",
        "Tabular Report": "You are an expert AI who can create a structured report from user request.The report should be in markdown format structured into subtopics/tables/lists. Do not add any additional comments.",
        "Tables only": "You are an expert AI who can create a structured tabular report from user request.The report should be in markdown format consists of only markdown tables. Do not add any additional comments.",
    },
    "online": {
        "Chat": "You are an expert AI, complete the given task using the provided context. Do not add any additional comments.",
        "Full Text Report": "You are an expert AI who can create a detailed report using information provided in the context from user request. The report should be in markdown format. Do not add any additional comments.",
        "Tabular Report": "You are an expert AI who can create a structured report using information provided in the context from user request. The report should be in markdown format structured into subtopics/tables/lists. Do not add any additional comments.",
        "Tables only": "You are an expert AI who can create a structured tabular report using information provided in the context from user request. The report should be in markdown format consists of only markdown tables. Do not add any additional comments.",
    },
}

class QueryModel(BaseModel):
    query: str = Query(default="market research", description="input query to generate Report")
    description: str = Query(default="", description="additional context for report")
    user_id: str = Query(default="", description="unique user id")
    user_name: str = Query(default="", description="user name")
    internet: bool = Query(default=True, description="Enable Internet search")
    output_format: str = Query(default="Tabular Report", description="Output format for the report", enum=["Chat", "Full Text Report", "Tabular Report", "Tables only"])
    data_format: str = Query(default="Structured data", description="Type of data to extract from the internet", enum=["No presets", "Structured data", "Quantitative data"])

@app.post("/generate_report")
async def generate_report(request: Request, query: QueryModel):
    query_str = query.query
    description = query.description
    user_id = query.user_id
    internet = "online" if query.internet else "offline"
    sys_prompt_output_format = sys_prompts[internet][query.output_format]
    data_format = query.data_format

    # Combine query with user keywords
    if query.internet:
        search_query = query_str 
        
        # Search for relevant URLs
        urls = search_brave(search_query, num_results=4)

        # Fetch and extract content from the URLs
        all_text_with_urls = fetch_and_extract_content(data_format, urls, query_str)

        # Prepare the prompt for generating the report
        additional_context = limit_tokens(str(all_text_with_urls))
        prompt = f"#### COMPLETE THE TASK: {query_str} #### IN THE CONTEXT OF ### CONTEXT: {description} #### ADDITIONAL CONTEXT:{additional_context}"

    else:
        prompt = f"#### COMPLETE THE TASK: {query_str} #### IN THE CONTEXT OF ### CONTEXT: {description}"
        all_text_with_urls = [("","")]

    md_report = together_response(prompt, model=llm_default_medium, SysPrompt=sys_prompt_output_format)
    
    # Insert data into database (or other storage)
    insert_data(user_id, query_str, description,  str(all_text_with_urls), md_report)
    references_html = dict()
    for text, url in all_text_with_urls:
        references_html[url] = str(md_to_html(text))


    # Return the generated report
    return {
        "report": md_to_html(md_report),
        "references": references_html
    }

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],)