main.py

import torch
import torch.serialization        # the module that controls loading behaviour
from supar.config import Config
import supar.utils.data as _supar_data
import multiprocess as mp
from transformers import BertTokenizerFast


# Patch the tokenizer class directly
if not hasattr(BertTokenizerFast, 'split_special_tokens'):
    BertTokenizerFast.split_special_tokens = False

# Force supar's pool to inherit the patch by using fork start method
mp.set_start_method('fork', force=True)

"""
main.py — Sebayhi FastAPI backend

Endpoints:
  GET  /              → serves the chat UI
  POST /chat/stream   → SSE streaming response

Intent is now determined by the frontend (user explicitly selects
"knowledge" or "i3rab" before typing). The backend receives it directly
in the request body — no classifier needed.
"""
import re
import asyncio
from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel
from typing import Literal
from haystack_pipeline import *
from pymongo import MongoClient
import json
import os
import requests 
import uuid
from dotenv import load_dotenv
from datetime import datetime
from fastapi.middleware.cors import CORSMiddleware
load_dotenv()

try:
    connection_uri = os.getenv("MONGODB_URL")
    client = MongoClient(connection_uri, tls=True) # temp fix

except Exception as e:
    print("An error occured when trying to connect to the cluster: ", e)

db = client.get_database("Sebayhi_convos")
collection = db.get_collection("collection_0")

# retrieve docs from mongodb and stringify them

    

# modify this function later
#def retrieve_docs(session_id: str):
#    #session_id = request.state.session_id
#    docs = list(collection.find({
#        "session_id": session_id
#    }))
#
#    return json.dumps(docs, indent=2, default=str)
    


def run_query(message: str, intent: str, session_id: str) -> str:
    """
    Routes the user message to the appropriate Haystack pipeline
    based on the intent selected in the UI.
 
    Args:
        message: raw Arabic text from the user
        intent:  "knowledge" | "i3rab"
 
    Returns:
        LLM reply as a string
    """
    #session_id = request.state.session_id
    match intent:
        case "i3rab":
            # Strip any leading instruction words so the parser gets
            # only the sentence: "أعرب: ذهب الطالب" → "ذهب الطالب"
            sentence = re.sub(
                r"^(أعرب|اعرب|حلل|إعراب|اعراب)\s*[:\-]?\s*",
                "",
                message.strip()
            )
            #i3rab_pipeline.run({"parser": {"sentence": sentence}})

            parser_output = get_camel_parser_output(sentence)

            if parser_output:
                print("\n--- تم استخراج البيانات بنجاح، جاري المعالجة... ---")
                final_tags_meaning = get_parsed_lines(parser_output)
                print("\n--- tags description. ---")
                print(final_tags_meaning)
                print("\n--- تم تفسير الوسوم بنجاح ---")
                print("\n--- جاري توليد الإعراب... ---")
                result = i3rab_final_output(final_tags_meaning)
                print("\n=== النتيجة النهائية ===")
                print(result)

                # test to see to insert/update msgs into mongodb docs
                #query_filter = {"role": "ai-assistant"}
                #update_operation = {"$push": {
                #    "message": result
                #}}
                #result_mongodb = collection.update_one(query_filter, update_operation)
                return result


            else:
                err_msg = "فشل المحلل في استخراج بيانات."
                print(err_msg)
                return err_msg

        
        
        case "knowledge":

            chat_docs = collection.find_one({
                "_id": session_id
            })
            msg_history = json.dumps(chat_docs["messages"][-5:], indent=2, default=str)
            
            result = knowledge_pipeline.run({
                "text_embedder": {"text": message},
                "prompt_builder": {
                    "question": message,
                    "history": msg_history
                },
            })
 
            replies = result.get("llm", {}).get("replies", [])
            if replies:
                # test to see to insert/update msgs into mongodb docs
                
                return replies[0].text if hasattr(replies[0], "text") else str(replies[0])
            
            return "[لم يتم الحصول على إجابة]"

# ─────────────────────────────────────────────────────────────────────────────

app = FastAPI(title="Sebayhi", description="Arabic Grammar Education Tutor")
app.mount("/static", StaticFiles(directory="static"), name="static")

app.add_middleware(
    CORSMiddleware,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)
# fastapi middleware here --> cookies / session generation logic
@app.middleware("http")
async def user_session_management(request: Request, call_next):
    
    print("\n==============================")
    print("REQUEST:", request.method, request.url.path)
    print("INCOMING COOKIES:", request.cookies)

    session_id = request.cookies.get("session_id")
    print("COOKIE:", session_id)
    session = None

    # this if statement is not running OR session_id is never saved and found on mongodb
    if session_id:
        print("COOKIE SESSION ID:", session_id)

        retrieved_session = collection.find_one({"_id": session_id})
        print("MONGO FOUND:", retrieved_session is not None)

        print("FOUND SESSION:", session_id is not None)
        if retrieved_session:
            collection.update_one(
                {"_id": session_id},
                {"$set": {"last_active": datetime.utcnow()}}
            )
            #session = retrieved_session
            session_id = retrieved_session["_id"]
        else:
            print("COOKIE EXISTS BUT SESSION DOES NOT")
            session_id = None
    
    if session_id is None:
        try:
            session_id = str(uuid.uuid4())
            print("CREATING NEW SESSION:", session_id)
            session = {
                "_id": session_id,
                "created_at": datetime.utcnow(),
                "last_active": datetime.utcnow(),
                "messages": []
            }
            usr_doc = collection.insert_one(session)
            

        except Exception as e:
            print("An error occured when trying to generate a session id / create new session: ", e) 
    
    request.state.session_id = session_id
    print("FINAL REQUEST SESSION:", request.state.session_id)
    # forward request to actual endpoint and get the response
    response = await call_next(request)

    
    response.set_cookie(
        key="session_id",
        value=session_id,
        max_age=60*60*24*7,
        httponly=True,
        samesite="lax",
        secure=True,
        path="/"
    )
    print("SETTING COOKIE:", session_id)
    print("==============================\n")
    return response
    

# ── Request model ─────────────────────────────────────────────────────────────

class ChatRequest(BaseModel):
    message: str
    intent: Literal["knowledge", "i3rab"]   # sent explicitly by the frontend


# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/")
async def serve_ui():
    
    return FileResponse("static/index.html")



@app.post("/chat/stream")
async def chat_stream(req: ChatRequest, request: Request):
    """
    SSE streaming endpoint.
    Runs the pipeline in a thread (blocking I/O), then streams
    the reply word-by-word to the frontend.
 
    Chunk protocol:
      data: <token>\n\n   — text to append
      data: [DONE]\n\n    — end of stream
    """
    

    async def token_generator():
        # Run the blocking pipeline call in a thread pool so FastAPI
        # stays non-blocking for other concurrent requests
        loop = asyncio.get_event_loop()
        full_reply = await loop.run_in_executor(
            None, run_query, req.message, req.intent, request.state.session_id
        )
        # via mongodb api, insert the msgs and replies  into the docs
        # then in haystack_pipeline.py retrieve the docs and inject them into the prompt
        
        # append 
        #query_filter1 = {"role": "user"}
        #query_filter2 = {"role": "ai-assistant"}
        #
        #update_operation1 = {"$push": {
        #    "message": req.message
        #}}
#
        #update_operation2 = {"$push": {
        #    "message": full_reply
        #}}
        
        #user_msgs = collection.update_one(query_filter1, update_operation1)
        #ai_msgs = collection.update_one(query_filter2, update_operation2)

        
        all_msgs = collection.update_one(
            {"_id": request.state.session_id},
            {"$push": {
                "messages": {"$each": [req.message, full_reply]}
            }}
        )
        
        for word in full_reply.split(" "):
            yield f"data: {word} \n\n"
            await asyncio.sleep(0.02)
 
        yield "data: [DONE]\n\n"
 
    return StreamingResponse(
        token_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )
 
# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)


'''
mongodb methods;
- insertOne, insertMany()
- find(), sort(), limit()
- include document body as a parameter or part of it within a method
- updateOne(), updateMany()
-$set, $unset for updating documents with key value pairs
'''