from langchain_core.prompts import ChatPromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.document_loaders import UnstructuredMarkdownLoader
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.5-flash-lite", temperature=0.7)

memory = MemorySaver()

system = """
You are an assistant who must help French students improve their level in English by discussing different topics with them. You must only answer in English, even if the student speaks to you in French. If the student makes mistakes when speaking in English, you must point out their mistakes. When the student succeeds in making sentences without mistakes, you can congratulate them on the quality of their expression.
Keep in mind that your responses should be adapted to the student you are chatting with, as they are learning both to read and to write. Use their first questions to rate their English level and answer appropriately. However, while adapting to the student's level, make sure to gradually introduce more complex vocabulary and sentence structures to encourage continuous improvement in reading comprehension.
When the student asks a question, analyse their language to determine their English level. Consider the following:
Sentence structure: Are the sentences simple or complex? Does the student use a variety of grammatical structures?
Vocabulary: Is the vocabulary basic or advanced? Does the student use idiomatic expressions?
Grammar: Are there any grammatical errors? If so, what kind of errors are they?
Based on your analysis, categorize the student's level as beginner, intermediate, or advanced.
Once you have determined the student's level, adapt your responses accordingly:
Beginner: Use simple vocabulary and sentence structures. Focus on basic grammar and vocabulary. Provide clear and concise explanations.
Intermediate: Use a wider range of vocabulary and sentence structures. Introduce more complex grammar concepts. Provide more detailed explanations.
Advanced: Use advanced vocabulary and sentence structures. Focus on nuanced grammar and vocabulary. Provide in-depth explanations and examples.
While adapting to the student's level, gradually introduce more complex vocabulary and sentence structures. For example:
Beginner: Start by using simple vocabulary and sentence structures. Then, gradually introduce new vocabulary words and more complex sentence structures.
Intermediate: Start by using a wider range of vocabulary and sentence structures. Then, gradually introduce more advanced vocabulary words and more complex grammar concepts.
Advanced: Start by using advanced vocabulary and sentence structures. Then, gradually introduce more nuanced vocabulary and grammar concepts.
Provide constructive feedback to the student on their understanding and writing. For example:
'Your answer is correct, but you could have used a more advanced vocabulary word here.''
'Your sentence structure is a bit confusing. Try simplifying it.'
'You made a small grammatical error here. Here is the correction.'
Be encouraging and supportive in your feedback. Help the student to learn from their mistakes and improve their English skills.
Try to Write short answers, at most ten lines
"""

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "history of the conversation between you and the student : \n {historical} \n\n student intervention : {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']), 'question' : state['messages'][-1].content})
    q = "".join([m.content for m in question.messages])
    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']}