from langchain_core.prompts import ChatPromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.document_loaders import DirectoryLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_huggingface import HuggingFaceEmbeddings
from typing import List
from langchain_chroma import Chroma
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 langgraph.checkpoint.memory import MemorySaver
from fastapi import FastAPI, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional
from PIL import Image
import base64
from io import BytesIO
import os 
import logging
import sys

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

app = FastAPI()

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


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

memory = MemorySaver()

glob_pattern="./*.md"
directory_path = "./documents"
loader = DirectoryLoader(directory_path, glob=glob_pattern)
documents = loader.load()

file_path = "./documents/prog.pdf"
loader = PyPDFLoader(file_path)
prog = loader.load()

system = """
Tu es un assistant spécialisé dans l'enseignement de l'enseigement Sciences Numériques et Technologie (SNT) de la classe de seconde'
Ton interlocuteur est un élève qui suit l'enseigement SNT en classe de seconde
Ton unique thème de conservation doit être l'enseignement SNT. Tu ne dois pas aborder d'autres thèmes que l'enseignement SNT
Tu ne dois pas faire d'erreur, répond à la question uniquement si tu es sûr de ta réponse
si tu ne trouves pas la réponse à une question, tu réponds que tu ne connais pas la réponse et que l'élève doit s'adresser à son professeur pour obtenir cette réponse
si l'élève n'arrive pas à trouver la réponse à un exercice, tu ne dois pas lui donner tout de suite la réponse, mais seulement lui donner des indications pour lui permettre de trouver la réponse par lui même
Tu dois uniquement répondre en langue française
Tu ne dois pas commencer tes interventions par "Assistant".
Tu trouveras ci-dessous un résumé du cours de SNT , tu devras veiller à ce que tes réponses ne sortent pas du cadre du programme de SNT
Si la question posée ne rentre pas dans le cadre du programme de SNT tu peux tout de même répondre en précisant bien que cette notion est hors programme
si tu proposes un exercice, tu dois bien vérifier que toutes les notions nécessaires à la résolution de l'exercice sont explicitement au programme de SNT
"""

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "Programme officiel de SNT: \n {prog} \n\n" "Résumé du cours de SNT : \n {documents} \n\n Historique conversation entre l'assistant et l'élève : \n {historical} \n\n Intervention de l'élève : {question}"),
    ]
)

def format_historical(hist):
    historical = []
    for i in range(0,len(hist)-2,2):
        historical.append("Utilisateur : "+hist[i].content[0]['text'])
        historical.append("Assistant : "+hist[i+1].content[0]['text'])
    return "\n".join(historical[-10:])


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

def chatbot(state : GraphState):
    question = prompt.invoke({'historical': format_historical(state['messages']),'prog':prog, 'documents':documents ,  'question' : state['messages'][-1].content[0]['text']})
    q = question.messages[0].content + question.messages[1].content
    if len(state['messages'][-1].content) > 1 :
        response = llm.invoke([HumanMessage(
            content=[
                {"type": "text", "text": q},
                state['messages'][-1].content[1]
            ])])
    else :
        response = llm.invoke([HumanMessage(
            content=[
                {"type": "text", "text": q}
            ])])
    return {"messages": [AIMessage(content=[{'type': 'text', 'text': response.content}])]}

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

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

app_chatbot = workflow.compile(checkpointer=memory)

@app.post('/request')
def request(id:Annotated[str, Form()], query:Annotated[str, Form()], image:Optional[UploadFile] = None):
    config = {"configurable": {"thread_id": id}}
    if image:
        try:
            img = Image.open(image.file)
            img_buffer = BytesIO()
            img.save(img_buffer, format='PNG')
            byte_data = img_buffer.getvalue()
            base64_img = base64.b64encode(byte_data).decode("utf-8")
            message = HumanMessage(
            content=[
                {'type': 'text', 'text': query},
                {'type': 'image_url', 'image_url': {"url": f"data:image/jpeg;base64,{base64_img}"}}
            ])
        except:
            return {"response":"Attention, vous m'avez fourni autre chose qu'une image. Renouvelez votre demande avec une image."}
        rep = app_chatbot.invoke({"messages": message},config, stream_mode="values")
    else :
        rep = app_chatbot.invoke({"messages": [HumanMessage(content=[{'type': 'text', 'text': query}])]},config, stream_mode="values")
    return {"response":rep['messages'][-1].content[0]['text']}