import json
import os
import uuid
import tempfile
import numpy as np
from PIL import Image
from deep_translator import GoogleTranslator
from ultralytics import YOLO
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from gtts import gTTS
import whisper as whisper_lib
import re
from rapidfuzz import fuzz

# ==============================
# Lazy Loaded Models (IMPORTANT)
# ==============================

embed_model = None
yolo_model = None
whisper_model = None

_model = None
_tokenizer = None

_index = None
_artifact_docs = None
_artifact_texts = None

# ==============================
# Lazy Getters
# ==============================

def get_embed_model():
    global embed_model
    if embed_model is None:
        embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    return embed_model


def get_yolo_model():
    global yolo_model
    if yolo_model is None:
        model_path = os.path.join("models", "best_egypt.pt")
        yolo_model = YOLO(model_path)
    return yolo_model


def get_whisper_model():
    global whisper_model
    if whisper_model is None:
        whisper_model = whisper_lib.load_model("small")
    return whisper_model


def get_llm():
    global _model, _tokenizer
    if _model is None:
        print("Loading LLM...")
        model_name = "google/flan-t5-base"
        _tokenizer = AutoTokenizer.from_pretrained(model_name)
        _model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        print("LLM loaded.")
    return _model, _tokenizer


# ==============================
# Load Artifacts JSON (lazy index)
# ==============================

import os
import json

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(BASE_DIR, "data", "artifacts.json")

with open(file_path, "r", encoding="utf-8") as f:
    data = json.load(f)


def build_index():
    global _index, _artifact_docs, _artifact_texts

    if _index is not None:
        return _index

    model = get_embed_model()

    artifact_texts = []
    artifact_docs = []

    for doc in data:
        name = doc.get("name", "")
        keywords = " ".join(doc.get("keywords", []))
        text = f"{name} {keywords}"

        artifact_texts.append(text)
        artifact_docs.append(doc)

    embeddings = model.encode(artifact_texts, convert_to_numpy=True)

    import faiss
    _index = faiss.IndexFlatL2(embeddings.shape[1])
    _index.add(np.array(embeddings))

    _artifact_docs = artifact_docs
    _artifact_texts = artifact_texts

    return _index


# ==============================
# Helper Functions
# ==============================

def is_arabic(text):
    return any('\u0600' <= c <= '\u06FF' for c in text)


def translate_to_en(text):
    try:
        return GoogleTranslator(source='auto', target='en').translate(text)
    except:
        return text


def translate_to_ar(text):
    try:
        return GoogleTranslator(source='auto', target='ar').translate(text)
    except:
        return text


# ==============================
# Intent Detection
# ==============================

INTENT_KEYWORDS = {
    "creator": ["who built", "who made", "creator", "made by",
                "من بناه", "من صنعه", "المنشئ", "من بنى"],
    "built_year": ["when was it built", "year", "date",
                   "متى بني", "سنة", "تاريخ"],
    "type": ["type", "what kind", "نوع", "ما نوع"],
    "era": ["era", "period", "dynasty",
            "العصر", "الفترة", "الحقبة", "عصر"],
    "material": ["material", "made of",
                 "مم صنع", "مصنوع من", "المادة"],
    "description": ["describe", "appearance", "look like",
                    "وصف", "كيف يبدو", "شكل"],
    "importance": ["importance", "significance", "why important",
                   "الأهمية", "أهميته", "ليه مهم"],
    "location_found": ["where was it found", "discovered",
                       "اين وجد", "مكان اكتشافه", "اكتشف"],
    "current_location": ["where is", "current location", "located",
                         "اين يوجد", "يقع", "مكانه"],
    "summary": ["tell me about", "overview", "summary", "brief",
                "what is", "who is", "information",
                "احكيلي", "أهم المعلومات", "نبذة", "معلومات"],
}


def detect_intents(q_en, q_ar=""):
    q_en = (q_en or "").lower()
    q_ar = (q_ar or "").lower()
    detected_intents = []

    for intent, keywords in INTENT_KEYWORDS.items():
        if intent == "summary":
            continue

        for kw in keywords:
            if kw in q_en or kw in q_ar:
                detected_intents.append(intent)
                break

    if not detected_intents:
        detected_intents.append("summary")

    return detected_intents


# ==============================
# Response Generator
# ==============================

def generate_intent_response(artifact, intent, user_lang="en"):
    name_en = artifact.get("name", "This artifact").replace("-", " ").replace("_", " ")
    name = translate_to_ar(name_en) if user_lang == "ar" else name_en

    value = artifact.get(intent, "Unknown")

    if user_lang == "ar":
        creator = translate_to_ar(str(artifact.get('creator', '')))
        era = translate_to_ar(str(artifact.get('era', '')))
        location = translate_to_ar(str(artifact.get('current_location', '')))
        description = translate_to_ar(str(artifact.get('description', '')))
        material = translate_to_ar(str(artifact.get('material', '')))
        value = translate_to_ar(str(value))

        templates = {
            "creator": f"تم إنشاء {name} بواسطة {creator}.",
            "built_year": f"تم بناء {name} في عام {value}.",
            "type": f"{name} هو {value}.",
            "era": f"يرجع {name} إلى عصر {era}.",
            "material": f"{name} مصنوع من {material}.",
            "description": f"يتميز {name} بأنه {description}، ويعكس أهمية كبيرة في تاريخ وحضارة مصر القديمة.",
            "importance": f"تكمن أهمية {name} في أنه {value}.",
            "location_found": f"تم اكتشاف {name} في {value}.",
            "current_location": f"يوجد {name} حاليًا في {location}.",
            "summary": f"يُعد {name} من أبرز المعالم الأثرية في مصر القديمة، حيث يتميز بأنه {description}. تم إنشاؤه بواسطة {creator}، ويرجع تاريخه إلى عصر {era}. ويقع حاليًا في {location}.",
        }

    else:
        templates = {
            "creator": f"{name} was created by {value}.",
            "built_year": f"{name} was built around {value}.",
            "type": f"{name} is a {value}.",
            "era": f"{name} dates back to the {value}.",
            "material": f"{name} is made of {value}.",
            "description": f"{name} is characterized by {value}.",
            "importance": f"The importance of {name}: {value}.",
            "location_found": f"{name} was discovered in {value}.",
            "current_location": f"{name} is currently located in {value}.",
            "summary": f"{name} is one of the most significant Egyptian monuments.",
        }

    return templates.get(intent, f"{name}: {value}")


# ==============================
# Artifact Search (lazy index)
# ==============================

def find_artifact(q_ar, q_en):
    model = get_embed_model()
    index = build_index()

    q_ar_n = q_ar.lower() if q_ar else ""
    q_en_n = q_en.lower() if q_en else ""

    query = q_en_n + " " + q_ar_n

    q_vec = model.encode([query])
    D, I = index.search(np.array(q_vec), k=5)

    candidates = [_artifact_docs[i] for i in I[0]]

    best_doc = None
    best_score = 0

    for doc in candidates:
        name = doc.get("name", "").lower()

        score = 0

        if name in query:
            score += 300

        score += fuzz.ratio(name, query)

        for token in name.split():
            if token in query:
                score += 50

        for kw in doc.get("keywords", []):
            if kw.lower() in query:
                score += 80

        if score > best_score:
            best_score = score
            best_doc = doc

    if best_score < 70:
        return None

    return best_doc


# ==============================
# YOLO (lazy)
# ==============================

def detect_artifact(image):
    model = get_yolo_model()

    results = model(image)
    result = results[0]

    annotated = Image.fromarray(result.plot())

    if result.boxes is None or len(result.boxes) == 0:
        return None, annotated

    best_idx = int(np.argmax(result.boxes.conf.cpu().numpy()))
    class_id = int(result.boxes.cls[best_idx])

    return result.names[class_id], annotated


# ==============================
# STT (lazy whisper)
# ==============================

def speech_to_text(audio_bytes):
    if not audio_bytes:
        return "", "en"

    tmp_path = None

    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
            f.write(audio_bytes)
            tmp_path = f.name

        model = get_whisper_model()

        audio = whisper_lib.load_audio(tmp_path)
        clip = whisper_lib.pad_or_trim(audio)
        mel = whisper_lib.log_mel_spectrogram(clip).to(model.device)

        _, probs = model.detect_language(mel)
        det_lang = max(probs, key=probs.get)

        res = model.transcribe(tmp_path, language=det_lang, fp16=False)
        text = res.get("text", "")

        return text, det_lang

    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)


# ==============================
# TTS
# ==============================

def text_to_speech(text):
    try:
        path = os.path.join(tempfile.gettempdir(), f"tts_{uuid.uuid4().hex}.mp3")
        gTTS(text=text, lang="ar" if is_arabic(text) else "en").save(path)
        return path
    except:
        return None


def cleanup_audio_file(filepath):
    if filepath and os.path.exists(filepath):
        os.remove(filepath)


# ==============================
# Chatbot Core (UNCHANGED LOGIC)
# ==============================

def chatbot_updated(question, image=None):

    if isinstance(question, bytes):
        text, lang = speech_to_text(question)
    else:
        text = (question or "").strip()
        lang = "ar" if is_arabic(text) else "en"

    if not text:
        return "Please provide a question." if lang == "en" else "من فضلك اكتب سؤالك."

    q_en = translate_to_en(text) if lang == "ar" else text

    artifact_from_image = None
    detected_name = None

    if image is not None:
        detected_name, _ = detect_artifact(image)
        if detected_name:
            artifact_from_image = find_artifact("", detected_name)

    artifact_from_text = find_artifact(text, q_en)

    if artifact_from_image:
        artifact = artifact_from_image if artifact_from_image else artifact_from_text
    else:
        artifact = artifact_from_text

    if not artifact:
        return "Artifact not found."

    intents = detect_intents(q_en, text)

    responses = []

    for intent in intents:
        responses.append(generate_intent_response(artifact, intent, lang))

    final_response = " ".join(responses)

    return final_response