import json
from openai import OpenAI
#from dotenv import load_dotenv
import os
import gradio as gr
from gradio import ChatMessage
import time
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

#load_dotenv()

#use openai API key
# api_key = os.getenv("OPENAI_API_KEY") 
# client = OpenAI(api_key=api_key)
#model="gpt-4.1"

#use gemini API key
api_key = os.getenv('GEMINI_API_KEY')

if api_key:
    logging.info("API Key Loaded Successfully!")
else:
    logging.info("API Key Missing!")


client = OpenAI(api_key=api_key,
         base_url="https://generativelanguage.googleapis.com/v1beta/openai/"
)
model="gemini-2.0-flash"


system_prompt = """
You are an AI assistant who is expert in breaking down complex problems into smaller, manageable parts.
Your task is to assist the user with their questions about Maths.

For the given user input, analyse the input and break down the problem step by step
Atleast 5-6 steps are required to solve the problem.

The steps are you get a user input, you analyse, you think, 
you think again for several times and then return an output 
with explanation and then finally you validate the output as well 
before returning the final output.

Rules:
1. Follow the strict JSON output as per output schema
2. Always perform one step at a time and wait for next input
3. Carefully analyse the user query
4. Makes sure not to give same answer in multiple steps.  Analyse the context to see what answers were provided earlier by the assistant and avoid repeating answers!  We need to make progress towards final answer.

Output format:
{{'step':"string", 'content':"string"}}

Example:
Input: What is 2 + 2.
Output: {{'step':"analyse",'content':"The user is asking for the sum of 2 and 2."}}
{{'step':"think",'content':"To find the sum, I need to add the two numbers together."}}
{{'step':"output",'content':"2 + 2 = 4."}}
{{'step':"validate",'content':"The answer is correct because 2 added to 2 equals 4."}}
{{'step':"result",'content':"2+2=4 and that is calculated by adding all the numbers"}}

"""
#Create a list of messages
#Using ChatML format with Gemini
history=[
    {"role": "system", "content": system_prompt}
]

def llm_response(message, history):
    if(len(history) ==0):
        history=[{"role": "system", "content": system_prompt}]
    history.append({"role": "user", "content": message})
    response = ChatMessage(content="", metadata={"title": "_Thinking_ step-by-step", "id": 0, "status": "pending"})
    yield response

    accumulated_thoughts = ""
    start_time = time.time()
    while True: 
        llm_response = client.chat.completions.create(
            model=model,
            response_format={"type": "json_object"},
            messages= history)

        parsed_response = json.loads(llm_response.choices[0].message.content)
        logging.info("________________________________")
        logging.info(parsed_response)
        logging.info("________________________________")
        time.sleep(5)
        history.append({"role": "assistant", "content": parsed_response['content']})
        if parsed_response['step'] == 'result':
            #We print how much time it took to get the final result
            response.metadata["status"] = "done"
            response.metadata["duration"] = time.time() - start_time
            yield response

            #And here is the final result
            thought = f"🤖Final Result: {parsed_response['content']}"
            response = [
			        response,
			        ChatMessage(
			            content=thought
			        )
			    ]
            
            yield response
            break
        else: #we have not reached final result yet
            
            thought = (parsed_response["step"], parsed_response["content"])            
            accumulated_thoughts += f"**{thought[0]}**: {thought[1]}\n"            
            response.content = accumulated_thoughts.strip()            
            yield response
 


demo = gr.ChatInterface(
    llm_response,
    title="Chain of Thought based LLM Chat Interface 🤔- amitguptaforwork@gmail.com",
    type="messages",
    theme='amitguptaforwork/blue_professional',
    examples=["how to calculate 2^2 + 5^2", "What is pythagorus theorem", "What is 2+5*6"],
    
)

demo.launch()