import os
from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from typing import List
from typing_extensions import TypedDict
from typing import Annotated
from langgraph.graph.message import AnyMessage, add_messages
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import END, StateGraph, START
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from database import create_table, return_title, return_exercise
import logging

create_table()

logger = logging.getLogger('uvicorn.error')
logger.setLevel(logging.DEBUG)

class Request(BaseModel):
    enonce : str
    code : str
    res_test : str

app = FastAPI()

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

load_dotenv()
os.environ["GROQ_API_KEY"] = os.getenv('GROQ_API_KEY')
os.environ["TOKENIZERS_PARALLELISM"] = "false"

llm = ChatGroq(model="llama-3.3-70b-versatile", temperature=0.5)

exercise = ""
system = """
Tu es un expert Python.
Tu dois aider un élève à résoudre un exercice de programmation Python.
Tu ne dois jamais donner la correction de l'exercice (même partiellement) à l'élève, juste lui donner des inddications lui permettant de résoudre lui même l'exercice
Tu dois t'adresser directement à l'élève.
L'élève ne peut pas te poser des questions, il peut juste te proposer son code.
Tu ne dois pas proposer à l'élève de te poser des questions
Il est inutile de proposer à l'élève de tester son code avec les exemples proposés.
Tu ne dois pas proposer aux élèves des modifications du programme qui sorte du cadre de l'exercice. Par exemple, pour l'exercice qui demande d'écrire une fonction moyenne, si dans l'énoncé il est précisé que l'on a un tableau non vide d'entier en paramètre, il est inutile de dire à l'élève que son programme doit gérer les tableaux vides.
Des tests unitaires ont été prévus afin d'évaluer le code proposé par l'élève. Si les tests sont tous positifs, cela signifie que le code proposé par l'élève est correct (tu dois alors lui dire que son code est correct et éventuellement lui donner des conseils pour encore améliorer la qualité de son code). Si les tests unitaires sont négatifs, tu dois aider l'élève à trouver ces erreurs.  
Tu dois t'exprimer en français
"""

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Voici l'exercice proposé à l'élève: \n {enonce} \n Voici le programme proposé par l'élève : {code} \n Voici le résultat des tests unitaires : {test_unit}"),
    ]
)

chain = prompt | llm | StrOutputParser()

class GraphState(TypedDict):
    enonce : str
    messages: Annotated[list[AnyMessage], add_messages]
    res_test : str
    

def chatbot(state : GraphState):
    msg_test = ""
    res_test =state['res_test'].content
    if res_test == '0' :
        msg_test = "Le code de l'élève n'a pas encore été testé avec les tests unitaires"
    elif res_test == '1':
        msg_test = "Le code de l'élève a bien été testé avec les tests unitaires, il a échoué à au moins un test unitaire"
    elif res_test == '2':
        msg_test = "Le code de l'élève a bien été testé avec les tests unitaires, il a réussi tous les tests unitaires"
    response = chain.invoke({'enonce': state['enonce'].content, 'code' : state['messages'][-1].content, 'test_unit' : msg_test} )
    return {"messages": [AIMessage(content=response)]}

workflow = StateGraph(GraphState)
workflow.add_node('chatbot', chatbot)

workflow.add_edge(START, 'chatbot')
workflow.add_edge('chatbot', END)

app_chatbot = workflow.compile()

@app.post('/request')
def request(req: Request):
    rep = app_chatbot.invoke({"enonce" : HumanMessage(content=req.enonce),"messages": [HumanMessage(content=req.code)], "res_test" : HumanMessage(content=req.res_test)}, stream_mode="values")
    return {"response":rep['messages'][-1].content}


@app.get('/title')
def get_title():
    tab_title = return_title()
    return {"title":tab_title}

@app.get('/exercise/{id}')
def get_exercise(id : int):
    ex = return_exercise(id)
    return {"title" : ex[1].replace("\n",""), "enonce" : ex[2], "test": ex[3]}