# ===============================
# 1️⃣ استيراد المكتبات
# ===============================
import os
os.environ["HF_HOME"] = "/tmp"
os.environ["TRANSFORMERS_CACHE"] = "/tmp"
os.environ["PEFT_HOME"] = "/tmp"

from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from peft import LoraConfig, get_peft_model, TaskType, PeftModel
from flask import Flask, request, jsonify
from flask_cors import CORS
import threading, torch, json
from datasets import Dataset
import difflib

# ===============================
# 2️⃣ إعداد المتغيرات
# ===============================
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"   # موديل أكبر وأكثر دقة
FLASK_PORT = 7860

SAVE_DIR = "/tmp/lora_model"
DATA_DIR = "/tmp/data"
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(SAVE_DIR, exist_ok=True)

DATASET_FILE = os.path.join(DATA_DIR, "dataset.json")

# ===============================
# 3️⃣ تحديد الجهاز
# ===============================
device = "cuda" if torch.cuda.is_available() else "cpu"
print("🔹 الجهاز المستخدم:", device)

# ===============================
# 4️⃣ تحميل الموديل الأساسي
# ===============================
print("🔄 جاري تحميل الموديل الأساسي...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

try:
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="auto" if device == "cuda" else {"": "cpu"},
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        trust_remote_code=True
    )
    model.to(device)
    print("✅ تم تحميل الموديل الأساسي بنجاح.")
except Exception as e:
    print("⚠️ فشل تحميل الموديل:", e)
    raise e

# ===============================
# 5️⃣ إعداد LoRA
# ===============================
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=4,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj"]
)
model = get_peft_model(model, lora_config)
model.to(device)
print("✅ تم تفعيل LoRA على الموديل.")

# ===============================
# 6️⃣ تحميل أوزان LoRA إذا وجدت
# ===============================
if os.path.exists(SAVE_DIR):
    try:
        base_model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            device_map="auto" if device == "cuda" else {"": "cpu"},
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            trust_remote_code=True
        )
        model = PeftModel.from_pretrained(base_model, SAVE_DIR)
        model.to(device)
        print("✅ تم تحميل أوزان LoRA المحفوظة.")
    except Exception as e:
        print("⚠️ لم يتم تحميل أوزان LoRA:", e)
else:
    print("ℹ️ لا يوجد أوزان محفوظة. سيتم استخدام الموديل الأساسي فقط.")

# ===============================
# 7️⃣ إعداد Flask API
# ===============================
app = Flask(__name__)
CORS(app)

trainer = None
current_dataset = None

# قاعدة بيانات تجريبية للمستخدمين
users_db = {
    "1": {"name": "أحمد السيد", "role": "طالب"},
    "2": {"name": "منى علي", "role": "محامية"},
    "3": {"name": "عمر حسن", "role": "عميل"}
}

# ===============================
# 7a️⃣ توليد النصوص
# ===============================
@app.route("/generate", methods=["POST"])
def generate():
    try:
        data = request.json or {}
        prompt = str(data.get("prompt", "")).strip()
        user_id = str(data.get("userId", "")).strip()
        focus_level = float(data.get("focus_level", 0.5))
        n_examples = int(data.get("n_examples", 2))

        user_info = users_db.get(user_id)
        user_name = user_info["name"] if user_info else "مستخدم"

        if not prompt:
            return jsonify({"messages": [{"role": "assistant", "content": "⚠️ من فضلك أرسل سؤال أو استفسار صحيح."}]}), 400

        # في حالة عدم وجود بيانات تدريبية
        if current_dataset is None:
            base_prompt = (
                f"أنت مساعد قانوني اسمه زين.\n"
                f"مهمتك: الرد على الأسئلة القانونية بلغة عربية واضحة وبسيطة "
                f"مع الحفاظ على الدقة القانونية.\n\n"
                f"السؤال:\n{prompt}\n\nالإجابة:"
            )

            inputs = tokenizer(base_prompt, return_tensors="pt", truncation=True, max_length=1024)
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
            input_len = inputs["input_ids"].shape[-1]

            outputs = model.generate(
                **inputs,
                max_new_tokens=300,
                temperature=0.6,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1
            )

            reply = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()

            return jsonify({
                "messages": [
                    {"role": "assistant", "content": f"مرحباً {user_name}"},
                    {"role": "assistant", "content": reply}
                ],
                "meta": {"mode": "base"}
            })

        # في حالة وجود بيانات تدريبية
        candidates = []
        for item in current_dataset:
            text = item.get("text", "")
            ratio = difflib.SequenceMatcher(None, prompt.lower(), text.lower()).ratio()
            candidates.append((ratio, text))

        candidates.sort(reverse=True, key=lambda x: x[0])
        top_texts = [t for r, t in candidates[:n_examples] if r > 0.05]

        context_parts = []
        for t in top_texts:
            lines = t.split("\n")
            u_lines, a_lines = [], []
            for l in lines:
                if ":" in l:
                    role, content = l.split(":", 1)
                    role = role.strip().lower()
                    if role.startswith("user"):
                        u_lines.append(content.strip())
                    elif role.startswith("assistant") or role.startswith("ai") or role.startswith("system"):
                        a_lines.append(content.strip())
            if u_lines or a_lines:
                context_parts.append("المستخدم: " + " ".join(u_lines) + "\nالمساعد: " + " ".join(a_lines))

        context_text = "\n\n---\n\n".join(context_parts)

        instruction = (
            "التعليمات: استخدم الأمثلة التدريبية كمصدر مرجعي، "
            "لكن أجب بصياغة جديدة وبأسلوبك الخاص. "
            "حافظ على صحة المعلومات القانونية وتجنب النسخ الحرفي."
        )

        if context_text and focus_level > 0.0:
            emphasis = "ركز بشكل كبير على الأمثلة" if focus_level > 0.7 else "خذ الأمثلة كمرجع فقط"
            final_prompt = (
                f"{instruction}\n\n{emphasis}\n\n"
                f"أمثلة تدريبية (مستوى التركيز={focus_level}):\n{context_text}\n\n"
                f"السؤال:\n{prompt}\n\nالإجابة:"
            )
        else:
            final_prompt = f"{instruction}\n\nالسؤال:\n{prompt}\n\nالإجابة:"

        # توليد النص
        inputs = tokenizer(final_prompt, return_tensors="pt", truncation=True, max_length=1024)
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        input_len = inputs["input_ids"].shape[-1]

        outputs = model.generate(
            **inputs,
            max_new_tokens=300,
            temperature=0.6,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.1
        )

        reply = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()

        return jsonify({
            "messages": [
                {"role": "assistant", "content": f"مرحباً {user_name}"},
                {"role": "assistant", "content": reply}
            ],
            "meta": {
                "mode": "trained",
                "used_examples": len(context_parts),
                "focus_level": focus_level
            }
        })

    except Exception as e:
        return jsonify({"messages": [{"role": "assistant", "content": f"خطأ: {str(e)}"}]})

# ===============================
# 7b️⃣ تحميل البيانات
# ===============================
def load_dataset():
    global current_dataset
    if os.path.exists(DATASET_FILE):
        with open(DATASET_FILE, "r", encoding="utf-8") as f:
            try:
                data = json.load(f)
            except:
                data = []
        if data:
            formatted = []
            for item in data:
                if "messages" in item:
                    conversation = ""
                    for msg in item["messages"]:
                        conversation += f"{msg['role'].capitalize()}: {msg['content']}\n"
                    formatted.append({"text": conversation.strip()})
                elif "text" in item:
                    formatted.append({"text": item["text"]})
            current_dataset = Dataset.from_list(formatted)
            print(f"✅ تم تحميل {len(formatted)} سجل من قاعدة البيانات.")
        else:
            current_dataset = None
    else:
        current_dataset = None

load_dataset()

# ===============================
# 7c️⃣ تدريب الموديل (LoRA)
# ===============================
@app.route("/train", methods=["POST"])
def train():
    global trainer, current_dataset, model
    try:
        data = []
        epochs = 1
        batch_size = 1

        if request.is_json:
            body = request.get_json()
            if isinstance(body, list):
                data = body
            elif isinstance(body, dict) and "data" in body:
                data = body["data"]
            else:
                return jsonify({"error": "صيغة JSON غير صحيحة"}), 400

            epochs = int(body.get("epochs", 1)) if isinstance(body, dict) else 1
            batch_size = int(body.get("batch_size", 1)) if isinstance(body, dict) else 1
        else:
            file = request.files.get("file")
            if file:
                data = json.load(file)
                epochs = int(request.form.get("epochs", 1))
                batch_size = int(request.form.get("batch_size", 1))
            else:
                return jsonify({"error": "لم يتم رفع ملف أو إرسال JSON"}), 400

        old_data = []
        if os.path.exists(DATASET_FILE):
            with open(DATASET_FILE, "r", encoding="utf-8") as f:
                try:
                    old_data = json.load(f)
                except:
                    old_data = []

        combined = old_data + data
        unique_data, seen_texts = [], set()
        for item in combined:
            if "messages" in item:
                text_key = "\n".join([f"{m['role'].strip().lower()}:{m['content'].strip()}" for m in item["messages"]])
            elif "text" in item:
                text_key = item["text"].strip()
            else:
                continue
            if text_key not in seen_texts:
                seen_texts.add(text_key)
                unique_data.append(item)

        with open(DATASET_FILE, "w", encoding="utf-8") as f:
            json.dump(unique_data, f, ensure_ascii=False, indent=2)

        formatted = []
        for item in unique_data:
            if "messages" in item:
                conversation = ""
                for msg in item["messages"]:
                    conversation += f"{msg['role'].capitalize()}: {msg['content']}\n"
                formatted.append({"text": conversation.strip()})
            elif "text" in item:
                formatted.append({"text": item["text"]})
        current_dataset = Dataset.from_list(formatted)

        def tokenize_function(examples):
            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256)

        tokenized_dataset = current_dataset.map(tokenize_function, batched=True, remove_columns=["text"])

        training_args = TrainingArguments(
            output_dir=SAVE_DIR,
            per_device_train_batch_size=batch_size,
            num_train_epochs=epochs,
            logging_steps=2,
            save_strategy="epoch",
            fp16=torch.cuda.is_available(),
            report_to=None
        )

        trainer = Trainer(model=model, args=training_args, train_dataset=tokenized_dataset)

        def run_training():
            print("🔄 بدأ التدريب...")
            trainer.train()
            trainer.save_model(SAVE_DIR)
            print("✅ انتهى التدريب وتم حفظ الموديل.")

        threading.Thread(target=run_training, daemon=True).start()

        return jsonify({
            "status": "بدأ التدريب",
            "epochs": epochs,
            "batch_size": batch_size,
            "records": len(unique_data)
        })

    except Exception as e:
        return jsonify({"error": str(e)})

# ===============================
# 8️⃣ نقطة البداية
# ===============================
@app.route("/", methods=["GET"])
def home():
    return jsonify({
        "status": "جاهز",
        "message": "واجهة برمجة التخرج تعمل الآن",
        "endpoints": ["/generate", "/train"]
    })

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=FLASK_PORT, threaded=True)