from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
import pandas as pd
import os
import requests
from io import StringIO
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import HfFolder
from tqdm import tqdm


app = FastAPI()

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  
    allow_credentials=True,
    allow_methods=["*"],  # Allow all methods
    allow_headers=["*"],  # Allow all headers
)

# Load your fine-tuned model and tokenizer
MODEL_NAME = "aubmindlab/aragpt2-medium"

try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
except Exception as e:
    raise RuntimeError(f"Failed to load model or tokenizer: {str(e)}")

# Define the general prompt template
general_prompt_template = """
أنت الآن نموذج لغة مخصص لتوليد نصوص عربية تعليمية بناءً على المادة والمستوى التعليمي. سيتم إعطاؤك مادة تعليمية ومستوى تعليمي، وعليك إنشاء نص مناسب بناءً على ذلك. النص يجب أن يكون:
1. واضحًا وسهل الفهم.
2. مناسبًا للمستوى التعليمي المحدد.
3. مرتبطًا بالمادة التعليمية المطلوبة.
4. قصيرًا ومباشرًا.

### أمثلة:
1. المادة: العلوم
   المستوى: الابتدائي
   النص: النباتات كائنات حية تحتاج إلى الماء والهواء وضوء الشمس لتنمو. بعض النباتات تنتج أزهارًا وفواكه. النباتات تساعدنا في الحصول على الأكسجين.

2. المادة: التاريخ
   المستوى: المتوسط
   النص: التاريخ هو دراسة الماضي وأحداثه المهمة. من خلال التاريخ، نتعلم عن الحضارات القديمة مثل الحضارة الفرعونية والحضارة الإسلامية. التاريخ يساعدنا على فهم تطور البشرية.

3. المادة: الجغرافيا
   المستوى: المتوسط
   النص: الجغرافيا هي دراسة الأرض وخصائصها. نتعلم عن القارات والمحيطات والجبال. الجغرافيا تساعدنا على فهم العالم من حولنا.

---

المادة: {المادة}
المستوى: {المستوى}

اكتب نصًا مناسبًا بناءً على المادة والمستوى المحددين. ركّز على جعل النص بسيطًا وواضحًا للمستوى المطلوب.
"""

class GenerateRequest(BaseModel):
    المادة: str
    المستوى: str

@app.post("/generate")
def generate_text(request: GenerateRequest):
    المادة = request.المادة
    المستوى = request.المستوى

    if not المادة or not المستوى or not isinstance(المادة, str) or not isinstance(المستوى, str):
        raise HTTPException(status_code=400, detail="المادة والمستوى مطلوبان ويجب أن يكونا نصًا.")

    try:
        # Format the prompt with user inputs
        arabic_prompt = general_prompt_template.format(المادة=المادة, المستوى=المستوى)

        # Tokenize the prompt
        inputs = tokenizer(arabic_prompt, return_tensors="pt", max_length=512, truncation=True)

        # Generate text
        with torch.no_grad():
            outputs = model.generate(
                inputs.input_ids,
                max_length=300,
                num_return_sequences=1,
                temperature=0.1,
                top_p=0.9,
                do_sample=True,
            )

        # Decode the generated text
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Remove the prompt from the generated text
        generated_text = generated_text.replace(arabic_prompt, "").strip()

        return {"generated_text": generated_text}

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during text generation: {str(e)}")

@app.get("/")
def read_root():
    return {"message": "Welcome to the Arabic Text Generation API!"}