# ===============================
# 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

# ===============================
# 2️⃣ إعداد المتغيرات
# ===============================
MODEL_NAME = "Qwen/Qwen2.5-0.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("🔹 Using device:", device)

# ===============================
# 4️⃣ تحميل الموديل الأساسي بأمان
# ===============================
print("🔄 Loading base model...")
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,
        load_in_8bit=(device=="cuda"),
        trust_remote_code=True
    )
    model.to(device)
    print("✅ Base model loaded!")
except Exception as e:
    print("⚠️ Failed to load base model:", e)
    raise e

# ===============================
# 5️⃣ إعداد LoRA
# ===============================
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=2,
    lora_alpha=8,
    lora_dropout=0.05,
    target_modules=["q_proj", "v_proj"]
)
model = get_peft_model(model, lora_config)
model.to(device)
print("✅ LoRA model ready on device:", device)

# ===============================
# 6️⃣ Auto load LoRA weights
# ===============================
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,
            load_in_8bit=(device=="cuda"),
            trust_remote_code=True
        )
        model = PeftModel.from_pretrained(base_model, SAVE_DIR)
        model.to(device)
        print("✅ LoRA weights auto-loaded from disk!")
    except Exception as e:
        print("⚠️ Failed to auto-load LoRA weights:", e)
else:
    print("ℹ️ No saved LoRA weights found, running with base model.")

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

trainer = None
current_dataset = None

# Fake DB
users_db = {
    "1": {"name": "Ahmed Elsayed", "role": "Student"},
    "2": {"name": "Mona Ali", "role": "Lawyer"},
    "3": {"name": "Omar Hassan", "role": "Client"}
}

# ===============================
# 7a️⃣ توليد نصوص
# ===============================
@app.route("/generate", methods=["POST"])
def generate():
    try:
        data = request.json
        prompt = data.get("prompt", "").strip()
        user_id = str(data.get("userId", "")).strip()

        user_info = users_db.get(user_id)
        if not user_info:
            return jsonify({"error": "User not found"}), 404

        user_name = user_info["name"]

        if current_dataset is None:
            return jsonify({
                "messages": [
                    {"role": "assistant", "content": f"⚖️ أهلاً {user_name}! أنا زين، مساعدك الذكي. برجاء طرح سؤال قانوني."}
                ]
            })

        # Search in dataset
        for item in current_dataset:
            lines = item["text"].split("\n")
            for line in lines:
                if line.lower().startswith("user:") and prompt in line:
                    response_messages = []
                    for l in lines:
                        if ":" in l:
                            role, content = l.split(":", 1)
                            response_messages.append({
                                "role": role.strip().lower(),
                                "content": content.strip()
                            })
                    response_messages.insert(0, {"role": "assistant", "content": f"أهلاً {user_name} 👋"})
                    return jsonify({"messages": response_messages})

        # Fallback: generate
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.7,
            do_sample=True
        )
        reply = tokenizer.decode(outputs[0], skip_special_tokens=True)

        return jsonify({
            "messages": [
                {"role": "assistant", "content": f"أهلاً {user_name} 👋"},
                {"role": "assistant", "content": reply}
            ]
        })

    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"✅ Loaded {len(formatted)} records from dataset.json")
        else:
            current_dataset = None
    else:
        current_dataset = None

load_dataset()

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

        # 1️⃣ JSON مباشر
        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 format not recognized"}), 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

        # 2️⃣ form-data ملف
        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": "No file uploaded or JSON provided"}), 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)

        # تجهيز Dataset للتدريب
        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)

        # تشغيل التدريب في Thread منفصل
        def run_training():
            print("🔄 Training started...")
            trainer.train()
            trainer.save_model(SAVE_DIR)
            print("✅ Training completed & model saved!")

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

        return jsonify({
            "status": "training started",
            "epochs": epochs,
            "batch_size": batch_size,
            "records": len(unique_data)
        })

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


# ===============================
# 8️⃣ تشغيل Flask
# ===============================
@app.route("/", methods=["GET"])
def home():
    return jsonify({
        "status": "ok",
        "message": "Graduation Project API is running 🚀",
        "endpoints": ["/generate", "/train"]
    })

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