from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
import torch, os

os.environ["HF_HOME"] = "/app/.cache"
os.environ["TRANSFORMERS_CACHE"] = "/app/.cache"
os.environ["HF_HUB_CACHE"] = "/app/.cache"

app = FastAPI(title="Hakayti Generator API")

API_KEY = os.getenv("API_KEY", "my_secret_key_123")
MODEL_ID = "Reem1/hakayti-merged"

bnb_cfg = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float32,  
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
)
print(" Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map={"": "cpu"},   
    quantization_config=bnb_cfg
)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
print(" Model loaded on:", next(model.parameters()).device)

few_shot_example = """<s>[User]: اسمي كريم وأحب الفضاء. 
اكتب لي قصة طويلة للأطفال (40 جملة على الأقل) مكتوبة بأسلوب قصصي ممتع.
[Assistant]:
كريم كان يحب النظر إلى السماء كل ليلة. كان يجلس في شرفة بيته، يراقب النجوم ويعدها واحدة تلو الأخرى. أحيانًا كان يتخيل أن النجوم مثل مصابيح صغيرة تضيء له الطريق نحو أحلامه الكبيرة.  
في المدرسة، لاحظ المعلم شغف كريم بالفضاء. أعطاه كتابًا مليئًا بالصور الملونة للكواكب والأقمار. انبهر كريم عندما رأى صور كوكب زحل وحلقاته اللامعة، وأخذ يتخيل نفسه يطير حوله.  
قرر كريم أن يصنع نموذجًا صغيرًا للمجموعة الشمسية. جمع كرات صغيرة وطلّاها بالألوان: الأزرق للأرض، الأحمر للمريخ، والأصفر للشمس. علّقها بخيوط في غرفته، فصارت كأنها سماء صغيرة يعيش فيها.  
في أحد الأيام، أعلنت المدرسة عن مسابقة للرسم بعنوان "رحلة إلى الفضاء". رسم كريم لوحة كبيرة يظهر فيها صاروخ ينطلق من الأرض، وحوله أطفال يلوّحون بأيديهم. فاز كريم بالجائزة الأولى، وكان سعيدًا جدًا لأن حلمه بالفضاء أصبح أقرب مما يتخيل.  
في المساء، جلس كريم مع والده على سطح المنزل. أشار والده إلى القمر وقال: "تخيّل يا كريم أن العلماء يدرسون الآن كيف يعيش الناس يومًا ما في الفضاء". ابتسم كريم وأجاب: "سأكون أنا واحدًا منهم".  
وهكذا تعلم كريم أن الحلم يبدأ بخطوة صغيرة، وأن حب المعرفة يمكن أن يفتح له أبوابًا واسعة نحو المستقبل.
</s>"""

def generate_story(name: str, hobby: str):
    prompt = f"""{few_shot_example}

[User]: اسمي {name} وأحب {hobby}. 
اكتب لي قصة طويلة للأطفال (40 جملة على الأقل) مكتوبة بنفس الأسلوب السابق،
وتأكد أن القصة كلها تدور حول {hobby} فقط.

[Assistant]:"""

    out = pipe(
        prompt,
        max_new_tokens=900,
        temperature=0.9,
        top_p=0.95,
        do_sample=True
    )

    full_text = out[0]["generated_text"]
    story = full_text.split("[Assistant]:")[-1].replace("</s>", "").strip()
    return story

class StoryRequest(BaseModel):
    name: str
    hobby: str

@app.post("/generate")
async def generate(req: StoryRequest, request: Request):
    header_key = request.headers.get("x-api-key") or request.headers.get("authorization")
    if API_KEY and header_key not in (API_KEY, f"Bearer {API_KEY}"):
        raise HTTPException(status_code=401, detail="Unauthorized")
    return {"story": generate_story(req.name, req.hobby)}