import os
import re
import json
import warnings
import gradio as gr
from gradio_client import Client
from huggingface_hub import HfApi
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ======================== كتم التحذيرات ========================
warnings.filterwarnings("ignore")

# ======================== الإعدادات العامة ========================
DEVICE = "cpu"
WHISPER_MODEL = "openai/whisper-base"
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

# المساحات الخارجية
RADAR_SPACE_URL = "Asem75/Aiocr_Radar"
VAULT_REPO_ID = "Asem75/My_teacher_volt"

# التوكنز من البيئة
HF_TOKEN = os.environ.get("HF_TOKEN", "").strip()
RADAR_API_KEY = os.environ.get("INTERNAL_API_KEY", "").strip()

# إنشاء عميل Hub للتعامل مع الخزنة
_hf_api = HfApi(token=HF_TOKEN) if HF_TOKEN else None

# ======================== الملف المرجعي للمواد المخصصة ========================
CUSTOM_SUBJECTS_PATH = "Curriculum_Core/custom_subjects.json"

def load_custom_subjects():
    """تحميل قائمة المواد المخصصة من الخزنة."""
    if not _hf_api:
        return []
    try:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(
            repo_id=VAULT_REPO_ID,
            filename=CUSTOM_SUBJECTS_PATH,
            repo_type="dataset",
            token=HF_TOKEN
        )
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data.get("subjects", [])
    except Exception:
        return []

def save_custom_subjects(subjects_list):
    """حفظ قائمة المواد المخصصة إلى الخزنة."""
    if not _hf_api:
        print("⚠️ لا يوجد HF_TOKEN")
        return
    try:
        content = json.dumps(
            {"subjects": subjects_list},
            ensure_ascii=False,
            indent=2
        )
        _hf_api.upload_file(
            path_or_fileobj=content.encode("utf-8"),
            path_in_repo=CUSTOM_SUBJECTS_PATH,
            repo_id=VAULT_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN
        )
        print(f"✅ تم حفظ المواد المخصصة: {subjects_list}")
    except Exception as e:
        print(f"⚠️ خطأ في حفظ المواد: {e}")

# ======================== المواد الأساسية ========================
BASE_SUBJECTS = [
    "الرياضيات",
    "العلوم",
    "اللغة العربية",
    "اللغة الإنجليزية",
    "التربية الإسلامية",
    "الفيزياء",
    "الكيمياء",
    "الأحياء",
    "التاريخ",
    "الجغرافيا",
    "التربية الوطنية",
    "الحاسوب",
    "الفلسفة",
    "الفرنسية"
]

def get_all_subjects():
    """دمج المواد الأساسية مع المخصصة."""
    custom = load_custom_subjects()
    all_subjects = list(dict.fromkeys(BASE_SUBJECTS + custom))
    return all_subjects

# ======================== إدارة التخصصات للتوجيهي ========================
CUSTOM_TRACKS_PATH = "Curriculum_Core/custom_tracks.json"

def load_custom_tracks():
    """تحميل قائمة التخصصات المخصصة للتوجيهي من الخزنة."""
    if not _hf_api:
        return []
    try:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(
            repo_id=VAULT_REPO_ID,
            filename=CUSTOM_TRACKS_PATH,
            repo_type="dataset",
            token=HF_TOKEN
        )
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data.get("tracks", [])
    except Exception:
        return []

def save_custom_tracks(tracks_list):
    """حفظ قائمة التخصصات المخصصة إلى الخزنة."""
    if not _hf_api:
        print("⚠️ لا يوجد HF_TOKEN")
        return
    try:
        content = json.dumps(
            {"tracks": tracks_list},
            ensure_ascii=False,
            indent=2
        )
        _hf_api.upload_file(
            path_or_fileobj=content.encode("utf-8"),
            path_in_repo=CUSTOM_TRACKS_PATH,
            repo_id=VAULT_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN
        )
        print(f"✅ تم حفظ التخصصات المخصصة: {tracks_list}")
    except Exception as e:
        print(f"⚠️ خطأ في حفظ التخصصات: {e}")

def get_all_tracks():
    """جلب جميع التخصصات (الفارغة افتراضياً)."""
    return load_custom_tracks()

# ======================== عدد الدروس ========================
MAX_LESSONS = 30

def get_lesson_numbers():
    """إرجاع قائمة بأرقام الدروس: الدرس 1، الدرس 2، ..."""
    return [f"الدرس {i}" for i in range(1, MAX_LESSONS + 1)]

# ======================== تحميل النموذج ========================
print("⏳ تحميل النموذج...")
_llm_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
_llm_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    device_map="cpu",
    low_cpu_mem_usage=True
)
print("✅ النموذج جاهز")

# ======================== Whisper ========================
_whisper_pipe = None

def load_whisper():
    global _whisper_pipe
    if _whisper_pipe is not None:
        return _whisper_pipe
    from transformers import pipeline
    _whisper_pipe = pipeline(
        "automatic-speech-recognition",
        model=WHISPER_MODEL,
        device=DEVICE,
        chunk_length_s=30
    )
    return _whisper_pipe

# ======================== تحويل رقم الدرس إلى صيغة الملف ========================
def lesson_to_file_prefix(lesson_str):
    """
    تحويل 'الدرس 1' إلى '01'، 'الدرس 15' إلى '15'.
    """
    match = re.search(r'\d+', lesson_str)
    if match:
        num = int(match.group())
        return f"{num:02d}"
    return "01"

# ======================== إدارة الخزنة السيادية ========================
def ensure_vault_structure(grade, semester, subject, lesson, user_name, track=None):
    """
    بناء المسار الشجري في الخزنة:
    - Curriculum_Core/{grade}/{semester}/{subject}/
    - البحث عن ملف يبدأ بـ {lesson_prefix}_ داخل المجلد.
    - Student_Hub/User_{user_name}/profile.json
    """
    if not _hf_api:
        print("⚠️ لا يوجد HF_TOKEN للتعامل مع الخزنة")
        return

    try:
        # مسار المناهج مع الفصل الدراسي
        curriculum_path = f"Curriculum_Core/{grade}/{semester}/{subject}/"
        lesson_prefix = lesson_to_file_prefix(lesson)
        print(f"📂 مسار المناهج: {curriculum_path}")
        print(f"🔍 البحث عن ملف يبدأ بـ: {lesson_prefix}_")

        # إنشاء ملف .gitkeep لحفظ المجلد
        _hf_api.upload_file(
            path_or_fileobj=b"",
            path_in_repo=f"{curriculum_path}.gitkeep",
            repo_id=VAULT_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN
        )

        # مسار الطالب
        student_path = f"Student_Hub/User_{user_name}/"
        profile = {
            "user_name": user_name,
            "grade": grade,
            "semester": semester,
            "subject": subject,
            "lesson": lesson,
            "track": track if track else None,
            "last_updated": None
        }

        _hf_api.upload_file(
            path_or_fileobj=json.dumps(
                profile,
                ensure_ascii=False,
                indent=2
            ).encode("utf-8"),
            path_in_repo=f"{student_path}profile.json",
            repo_id=VAULT_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN
        )
        print(f"👤 تم تحديث ملف الطالب: {student_path}profile.json")

    except Exception as e:
        print(f"⚠️ خطأ في الخزنة: {e}")

# ======================== الرادار الصوتي ========================
def call_radar_api(text, subject, detected_emotion):
    try:
        client = Client(RADAR_SPACE_URL)
        word_count = len(text.split())

        if word_count <= 35:
            speaker_name = "ar-EG-SalmaNeural"
        else:
            voice_map = {
                "الرياضيات": "ar-JO-TaimNeural",
                "اللغة العربية": "ar-SA-HamedNeural",
                "اللغة الإنجليزية": "en-US-JennyNeural",
                "التربية الإسلامية": "ar-SA-HamedNeural"
            }
            speaker_name = voice_map.get(subject, "ar-JO-TaimNeural")

        audio_url = client.predict(
            text,
            speaker_name,
            RADAR_API_KEY,
            api_name="/synth_arabic"
        )

        if audio_url and isinstance(audio_url, (list, tuple)):
            audio_url = audio_url[0]
        return str(audio_url).strip() if audio_url else None
    except Exception as e:
        print(f"⚠️ خطأ الرادار: {e}")
        return None

# ======================== دالة المحادثة المطورة ========================
def teacher_chat(
    audio_mic,
    subject,
    student_name,
    grade,
    semester,
    lesson,
    track,
    chat_history_state
):
    if audio_mic is None:
        return chat_history_state, None, "🎤 الرجاء تسجيل السؤال أولاً"

    # 1. تحويل الصوت لنص
    try:
        whisper = load_whisper()
        result = whisper(audio_mic, generate_kwargs={"language": "arabic"})
        user_text = result["text"].strip()
    except Exception as e:
        return chat_history_state, None, f"⚠️ خطأ: {e}"

    if not user_text:
        return chat_history_state, None, "❌ لم يتم التعرف على كلام"

    print(f"🗣️ الطالب {student_name}: {user_text}")

    # 2. بناء messages
    context_parts = [f"مادة {subject}", f"{lesson}"]
    if semester:
        context_parts.insert(1, f"الفصل {semester}")
    if track:
        context_parts.append(f"تخصص {track}")

    context_str = " - ".join(context_parts)

    system_prompt = (
        f"أنت الأستاذ عبود، معلم خبير. "
        f"السياق الحالي: {context_str}. "
        "أجب بإيجاز (30 كلمة كحد أقصى). "
        "لا تعط الإجابة كاملة، قسمها لخطوات. "
        "اطرح سؤالاً في النهاية. "
        "أضف [حالة: حماس] أو [حالة: هدوء]."
    )
    messages = [{"role": "system", "content": system_prompt}]
    for msg in chat_history_state:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": user_text})

    # 3. توليد الرد
    try:
        text_prompt = _llm_tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        model_inputs = _llm_tokenizer(
            [text_prompt],
            return_tensors="pt"
        ).to(DEVICE)

        with torch.no_grad():
            generated_ids = _llm_model.generate(
                **model_inputs,
                max_new_tokens=80,
                temperature=0.4,
                do_sample=True,
                pad_token_id=_llm_tokenizer.eos_token_id,
                use_cache=True
            )

        generated_ids = [
            output_ids[len(input_ids):]
            for input_ids, output_ids in zip(
                model_inputs.input_ids,
                generated_ids
            )
        ]
        full_response = _llm_tokenizer.batch_decode(
            generated_ids,
            skip_special_tokens=True
        )[0].strip()

        emotion_match = re.search(r'\[حالة:\s*(\w+)\]', full_response)
        detected_emotion = emotion_match.group(1) if emotion_match else "هدوء"
        cleaned_response = re.sub(
            r'\[حالة:\s*\w+\]',
            '',
            full_response
        ).strip()

        if not cleaned_response.endswith(('.', '؟', '!')):
            cleaned_response += '.'

        print(f"🤖 الأستاذ: {cleaned_response}")
    except Exception as e:
        return chat_history_state, None, f"⚠️ خطأ النموذج: {e}"

    # 4. تحديث الخزنة السحابية مع الفصل الدراسي والتخصص
    ensure_vault_structure(
        grade=grade,
        semester=semester,
        subject=subject,
        lesson=lesson,
        user_name=student_name,
        track=track
    )

    # 5. تحديث التاريخ
    chat_history_state = chat_history_state + [
        {"role": "user", "content": user_text},
        {"role": "assistant", "content": cleaned_response}
    ]

    # 6. توليد الصوت
    audio_path = call_radar_api(cleaned_response, subject, detected_emotion)
    status = f"✅ ({detected_emotion})" if audio_path else "⚠️ فشل الصوت"
    return chat_history_state, audio_path, status

# ======================== الواجهة المطورة ========================
GRADES = [
    "الروضة",
    "الصف الأول",
    "الصف الثاني",
    "الصف الثالث",
    "الصف الرابع",
    "الصف الخامس",
    "الصف السادس",
    "الصف السابع",
    "الصف الثامن",
    "الصف التاسع",
    "الصف العاشر",
    "الصف الحادي عشر",
    "الصف الثاني عشر (التوجيهي)"
]

SEMESTERS = ["الفصل الأول", "الفصل الثاني"]

ALL_SUBJECTS = get_all_subjects()
LESSON_NUMBERS = get_lesson_numbers()
ALL_TRACKS = get_all_tracks()

with gr.Blocks(title="الأستاذ عبود", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎓 المنصة التعليمية - الأستاذ عبود (النسخة المطورة)")

    chat_history_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=2):
            # حقول تتبع الطالب
            with gr.Row():
                student_name_input = gr.Textbox(
                    label="👤 اسم الطالب",
                    value="Divid",
                    placeholder="أدخل اسم الطالب"
                )
                grade_dropdown = gr.Dropdown(
                    choices=GRADES,
                    value="الصف العاشر",
                    label="📆 الصف الدراسي"
                )

            # الفصل الدراسي + التخصص للتوجيهي
            with gr.Row():
                semester_dropdown = gr.Dropdown(
                    choices=SEMESTERS,
                    value="الفصل الأول",
                    label="📅 الفصل الدراسي"
                )
                # حقل التخصص - يظهر فقط عند اختيار التوجيهي
                track_dropdown = gr.Dropdown(
                    choices=ALL_TRACKS,
                    value=None,
                    label="🎯 التخصص (للتوجيهي فقط)",
                    visible=False,
                    allow_custom_value=True
                )

            # صف إضافة تخصص جديد (يظهر فقط مع التوجيهي)
            with gr.Row(visible=False) as track_input_row:
                new_track_input = gr.Textbox(
                    label="➕ أضف تخصصاً جديداً",
                    placeholder="مثال: صناعي، زراعي، تجاري..."
                )
                add_track_btn = gr.Button(
                    "💾 حفظ التخصص",
                    variant="secondary",
                    size="sm"
                )

            with gr.Row():
                subject_dropdown = gr.Dropdown(
                    choices=ALL_SUBJECTS,
                    value=ALL_SUBJECTS[0] if ALL_SUBJECTS else "الرياضيات",
                    label="📚 المادة الدراسية",
                    allow_custom_value=True
                )
                lesson_dropdown = gr.Dropdown(
                    choices=LESSON_NUMBERS,
                    value="الدرس 1",
                    label="📖 رقم الدرس"
                )

            # إضافة مادة جديدة
            with gr.Row():
                new_subject_input = gr.Textbox(
                    label="➕ أضف مادة جديدة (مخصصة)",
                    placeholder="اكتب اسم المادة ثم اضغط الزر..."
                )
                add_subject_btn = gr.Button(
                    "💾 حفظ المادة في الخزنة",
                    variant="secondary",
                    size="sm"
                )

            # رسالة حالة الخزنة
            vault_status = gr.Textbox(
                label="📋 حالة الخزنة",
                value="🟢 جاهزة",
                interactive=False
            )

            # المحادثة
            chatbot_display = gr.Chatbot(label="💬 الحوار", height=400)
            mic_input = gr.Audio(
                label="🎤 سجل سؤالك",
                type="filepath",
                sources=["microphone"]
            )
            send_btn = gr.Button("🚀 إرسال السؤال", variant="primary")

        with gr.Column(scale=1):
            gr.Markdown("### 📊 لوحة المعلومات")
            gr.Markdown(
                f"""
                | الإعداد | القيمة |
                |---------|--------|
                | الخزنة | `{VAULT_REPO_ID}` |
                | الرادار | `{RADAR_SPACE_URL}` |
                | المواد الأساسية | {len(BASE_SUBJECTS)} مادة |
                | المواد المخصصة | {len(load_custom_subjects())} مادة |
                | الدروس المتاحة | 1 - {MAX_LESSONS} |
                """
            )
            gr.Markdown("---")
            audio_output = gr.Audio(
                label="🔊 رد الأستاذ",
                type="filepath",
                autoplay=True
            )
            status_output = gr.Textbox(
                label="📌 حالة الطلب",
                interactive=False
            )

    # ════════════════════ دوال الأحداث ════════════════════

    def toggle_track_fields(grade):
        """إظهار/إخفاء حقول التخصص عند اختيار التوجيهي."""
        if "التوجيهي" in grade:
            return (
                gr.update(visible=True),
                gr.update(visible=True),
                gr.update(visible=True)
            )
        else:
            return (
                gr.update(visible=False, value=None),
                gr.update(visible=False),
                gr.update(visible=False)
            )

    grade_dropdown.change(
        fn=toggle_track_fields,
        inputs=[grade_dropdown],
        outputs=[track_dropdown, track_input_row, add_track_btn]
    )

    def add_custom_subject(new_subject):
        """إضافة مادة جديدة إلى الخزنة وتحديث القائمة المنسدلة."""
        if not new_subject or not new_subject.strip():
            return gr.update(), "⚠️ الرجاء كتابة اسم المادة أولاً"

        new_subject = new_subject.strip()
        current_custom = load_custom_subjects()

        if new_subject in BASE_SUBJECTS:
            return (
                gr.update(),
                f"⚠️ '{new_subject}' موجودة مسبقاً في المواد الأساسية"
            )
        if new_subject in current_custom:
            return (
                gr.update(),
                f"⚠️ '{new_subject}' موجودة مسبقاً في المواد المخصصة"
            )

        current_custom.append(new_subject)
        save_custom_subjects(current_custom)

        updated_subjects = get_all_subjects()
        return (
            gr.update(choices=updated_subjects, value=new_subject),
            f"✅ تمت إضافة '{new_subject}' إلى الخزنة بنجاح!"
        )

    add_subject_btn.click(
        fn=add_custom_subject,
        inputs=[new_subject_input],
        outputs=[subject_dropdown, vault_status]
    )

    def add_custom_track(new_track):
        """إضافة تخصص جديد للتوجيهي."""
        if not new_track or not new_track.strip():
            return gr.update(), "⚠️ الرجاء كتابة اسم التخصص أولاً"

        new_track = new_track.strip()
        current_tracks = load_custom_tracks()

        if new_track in current_tracks:
            return (
                gr.update(),
                f"⚠️ '{new_track}' موجود مسبقاً في التخصصات"
            )

        current_tracks.append(new_track)
        save_custom_tracks(current_tracks)

        updated_tracks = get_all_tracks()
        return (
            gr.update(choices=updated_tracks, value=new_track),
            f"✅ تمت إضافة تخصص '{new_track}' بنجاح!"
        )

    add_track_btn.click(
        fn=add_custom_track,
        inputs=[new_track_input],
        outputs=[track_dropdown, vault_status]
    )

    def process_and_display(
        audio_mic,
        student_name,
        grade,
        semester,
        subject,
        lesson,
        track,
        history_state
    ):
        new_history, audio, status = teacher_chat(
            audio_mic=audio_mic,
            subject=subject,
            student_name=student_name,
            grade=grade,
            semester=semester,
            lesson=lesson,
            track=track,
            chat_history_state=history_state
        )
        return (
            new_history,
            new_history,
            audio,
            f"🟢 {status}" if audio else f"🔴 {status}",
            None
        )

    send_btn.click(
        fn=process_and_display,
        inputs=[
            mic_input,
            student_name_input,
            grade_dropdown,
            semester_dropdown,
            subject_dropdown,
            lesson_dropdown,
            track_dropdown,
            chat_history_state
        ],
        outputs=[
            chat_history_state,
            chatbot_display,
            audio_output,
            status_output,
            mic_input
        ]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)