import os
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langgraph.checkpoint.memory import MemorySaver
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_exercise
import logging
from langchain_google_genai import ChatGoogleGenerativeAI

create_table()

GOOGLE_API_KEY = ''

llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0.7)


prompt_aide = PromptTemplate.from_template(
"""
Tu es un expert en pédagogie de l'apprentissage de la programmation
Le langage utilisé pour l'apprentissage de la programmation est Python.
Tu dois aider un élève à résoudre une activité de programmation Python.
Il existe 2 types d'activités :
- les activités qui demandent à l'élève d'écrire un programme Python complet
- les activités qui demandent à l'élève de compléter un programme Python à trous (les trous sont indiqués par 3 points : ...)
Si l'activité est de type "à trous" et que l'élève à laissé des trous, tu dois lui proposer de compléter les trous.
Tu ne dois jamais donner la solution de l'activité (même partiellement) à l'élève, juste lui donner des inddications lui permettant de résoudre lui même l'activité
L'élève est un débutant en programmation, tu dois lui rappeler les fondamentaux si nécessaire.
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'activité. Par exemple, pour l'activité 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.
Tu dois t'exprimer en français
Voici l'énoncé de l'activité :
{enonce}
{code_trous}
Voici le programme proposé par l'élève pour résoudre l'activité :
{code}
Pour améliorer ta réponse, tu as aussi à ta disposition l'historique des différents programme proposés par l'élève et les différents conseils que tu lui a déjà donné :
{historique} 
""")

def history(hist):
    historical = ""
    for i in range(len(hist)):
        if i%2 == 0:
            historical += "code de l'éléve : \n"+hist[i].content+"\n"
        else :
            historical += "aide de l'expert : \n"+hist[i].content+"\n"
    return historical

memory = MemorySaver()
app = FastAPI()

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

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

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

def aide(state : AgentState):
    llm_aide = prompt_aide | llm | StrOutputParser()
    if state['code_trous'] == "":
        c_trous = ""
    else :
        c_trous = "Voici le code à trous de l'activité : \n"+state['code_trous']
    response = llm_aide.invoke({'enonce': state['enonce'], "code_trous" : c_trous, 'code' : state['messages'][-1].content, 'historique' : history(state['messages'])})
    return {"messages": [AIMessage(content=response)]}

workflow = StateGraph(AgentState)

workflow.add_node("aide", aide)

workflow.add_edge(START, "aide")
workflow.add_edge( "aide", END)

graph = workflow.compile(checkpointer=memory)

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


@app.get('/exercise/{titre}')
def get_exercise(titre : str):
    ex = return_exercise(titre)
    return {"title" : ex[1], "enonce" : ex[2], "code_trous":ex[3], "test": ex[4]}