from langchain_core.prompts import ChatPromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.document_loaders import DirectoryLoader
from langchain_community.document_loaders import PyPDFLoader
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 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)
doc_2025 = loader.load()

file_path_prob = "./documents/prob.pdf"
loader_prob = PyPDFLoader(file_path_prob)
prob = loader_prob.load()

file_path_guide = "./documents/guide.pdf"
loader_guide = PyPDFLoader(file_path_guide)
guide = loader_guide.load()

file_path_QR = "./documents/QR.pdf"
loader_QR = PyPDFLoader(file_path_QR)
QR = loader_QR.load()

system = """
Tu es un assistant expert en pédagogie. Ta spécialité est le **grand oral** (épreuve du bacalauréat technique et général).
Ton interlocuteur est un élève qui est en train de préparer son épreuve de grand oral.
Ton rôle est d'aider l'élève à préparer son grand oral
Tu dois répondre à toutes ses questions : organisation de l'épreuve, but de l'épreuve...
Tu dois aussi l'aider à trouver un sujet et à construire sa problématique.
Tu dois uniquement discuter du grand oral avec ton interlocuteur
Le projet d'orientation post baccalauréat ne doit plus être obligatoirement évoqué au cours de l'épreuve (même si tu trouves dans les documents fournis des éléments qui laisseraient penser le contraire)
Si tu ne connais pas la réponse à une question, propose à l'élève de demander à son professeur. 
"""

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", """
        Voici différents documents qui t'aideront à répondre aux questions des élèves :
        Le guide du grand oral :
        {guide}
        Un ensemble de Question-Réponse sur le Grand Oral:
        {QR}
        Un texte qui explique la notion de sujet et de problématique:
        {prob}
        Les informations à jour sur le déroulement de l'épreuve (si tu trouves des informations qui te semble contradictoires dans les documents ci-dessous, c'est le documents ci-dessous qui doit être ta source d'informations) :
        {doc_2025}
        Tu trouveras aussi l'historique conversation avec l'élève : \n {historical}
        Et enfin l'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[-20:])


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

def chatbot(state : GraphState):
    question = prompt.invoke({'historical': format_historical(state['messages']),'prob':prob, 'doc_2025':doc_2025, 'guide':guide, 'QR':QR ,  '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']}