from fuzzy_json import loads
from half_json.core import JSONFixer
from together import Together
from retry import retry
import re
from dotenv import load_dotenv
import os
from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware

# Retrieve environment variables
TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")

SysPromptDefault = "You are an expert AI, complete the given task. 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."

@retry(tries=3, delay=1)
def together_response(message, model = "meta-llama/Llama-3-8b-chat-hf", SysPrompt = SysPromptDefault):
  client = Together(api_key=TOGETHER_API_KEY)

  messages=[{"role": "system", "content": SysPrompt},{"role": "user", "content": message}]

  response = client.chat.completions.create(
      model=model,
      messages=messages,
      temperature=0.2,
  )
  return response.choices[0].message.content

def json_from_text(text):
    """
    Extracts JSON from text using regex and fuzzy JSON loading.
    """
    match = re.search(r'\{[\s\S]*\}', text)
    if match:
      json_out = match.group(0)
    else:
      json_out = text
    try:
        # Using fuzzy json loader
        return loads(json_out)
    except Exception:
        # Using JSON fixer/ Fixes even half json/ Remove if you need an exception
        fix_json = JSONFixer()
        return loads(fix_json.fix(json_out).line)
    
def generate_topics(user_input,num_topics,previous_query):
    prompt = f"""create a list of {num_topics} subtopics to follow for conducting {user_input} in the context of {previous_query}, RETURN VALID PYTHON LIST"""
    response_topics = together_response(prompt, model = "meta-llama/Llama-3-8b-chat-hf", SysPrompt = SysPromptList)
    subtopics = json_from_text(response_topics)
    return subtopics

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

# Define the app
app = FastAPI()

# Create a Pydantic model to handle the input data
class TopicInput(BaseModel):
    user_input: str
    num_topics: int
    previous_query: str
    

@app.get("/",tags=["Home"])
def api_home():
    return {'detail': 'Welcome to FastAPI Subtopics API! /n visit https://pvanand-generate-subtopics.hf.space/docs to test'}


@app.post("/generate_topics/")
async def create_topics(input: TopicInput):
    topics = generate_topics(input.user_input, input.num_topics, input.num_topics.previous_query)
    return {"topics": topics}