import torch
import faiss
import numpy as np
import os
import re
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List, Dict
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer

# --- 1. إعداد الـ API ---
app = FastAPI(title="Agricultural Expert API (Free Tier)")

# --- 2. تحميل الموديلات (نسخة خفيفة للـ CPU) ---
model_name = "Qwen/Qwen2.5-1.5B-Instruct" 

print("⏳ Loading Models on CPU (This might take a minute)...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="cpu"  # إجبار الموديل يشتغل على البروسيسور
)
embedding_model = SentenceTransformer("BAAI/bge-m3", device="cpu") # الـ Embedding كمان على الـ CPU

# --- 3. بناء الـ FAISS Index ---
DISEASES_FOLDER = "diseases"
all_chunks = []
all_embeddings = []

def split_text(text, chunk_size=600, chunk_overlap=100):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - chunk_overlap
    return chunks

if os.path.exists(DISEASES_FOLDER):
    for filename in os.listdir(DISEASES_FOLDER):
        if filename.endswith(".txt"):
            with open(os.path.join(DISEASES_FOLDER, filename), "r", encoding="utf-8") as f:
                text = f.read()
                chunks = split_text(text)
                embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
                all_chunks.extend(chunks)
                all_embeddings.append(embeddings)

    if all_embeddings:
        all_embeddings = np.vstack(all_embeddings)
        index = faiss.IndexFlatL2(all_embeddings.shape[1])
        index.add(all_embeddings.astype('float32'))
        print(f"✅ Index built successfully with {len(all_chunks)} chunks.")
else:
    print("⚠️ Folder 'diseases' not found!")

# --- 4. هيكل الـ Request اللي الفرونت إند هيبعته ---
class ChatRequest(BaseModel):
    question: str
    chat_history: List[Dict[str, str]] = [] # الفرونت إند هيبعت الذاكرة هنا
    k_results: int = 3

# --- 5. الـ Endpoint اللي هيستقبل الأسئلة ---
@app.post("/ask")
async def ask_ai_api(request: ChatRequest):
    # 1. البحث في السياق
    question_embedding = embedding_model.encode([request.question], convert_to_numpy=True)
    distances, indices = index.search(question_embedding.astype('float32'), request.k_results)
    context = "\n\n".join([all_chunks[i] for i in indices[0]])

    # 2. بناء الرسائل بنظام "مرآة اللغة"
    messages = [
        {
            "role": "system",
            "content": (
                "You are a professional agricultural assistant.\n"
                "STRICT RULE: Always respond in the SAME LANGUAGE as the user's current question.\n"
                "- If the user asks in English, your entire response MUST be in English.\n"
                "- If the user asks in Arabic, your entire response MUST be in Arabic.\n"
                "Use the provided context to answer accurately. Don't use Chinese."
            )
        }
    ]
    
    # إضافة الذاكرة الجاية من الفرونت إند
    messages.extend(request.chat_history)
    
    # إضافة السؤال الحالي مع السياق
    messages.append({
        "role": "user", 
        "content": f"Context: {context}\n\nQuestion: {request.question}\n(أجب بشرح مفصل وجمل واضحة)"
    })

    # 3. تجهيز الـ Prompt للموديل
    text_input = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text_input], return_tensors="pt").to(model.device)

    # 4. التوليد (بإعدادات تمنع التقطيع)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=800,      
            temperature=0.5,         
            top_p=0.92,
            repetition_penalty=1.05, 
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )

    # استخراج النص
    generated_ids = outputs[0][inputs.input_ids.shape[1]:]
    answer = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

    # فلتر اللغات الغريبة
    answer = re.sub(r'[^\u0600-\u06FF\u0660-\u0669a-zA-Z0-9\s\.,!\?؟\-\%\(\)\:\/\*]', '', answer)

    # الرد النهائي اللي هيرجع للفرونت إند كـ JSON
    return {
        "question": request.question,
        "answer": answer,
        "retrieved_context": context
    }