import os
import sys
import gc
from functools import lru_cache

import gradio as gr
import torch
import torch.nn as nn
import numpy as np
import spaces
import librosa
import torchaudio
import torchaudio.transforms as T

from PIL import Image
from transformers import (
    pipeline,
    AutoModelForAudioClassification,
    AutoFeatureExtractor,
)
from huggingface_hub import hf_hub_download
import importlib.util


# ============================================================
# CONFIGURATION
# ============================================================

_DEVICE = "cpu"

IMAGE_MODEL_REPO = "xRayon/convnext-ai-images-detector"
IMAGE_CHECKPOINT = "AI Images Detector/checkpoints/checkpoint_phase2.pth"
IMAGE_ARCHITECTURE = "AI Images Detector/model/convnext.py"
IMAGE_TRANSFORMS = "AI Images Detector/transforms.py"
TEXT_MODEL_ID = "Rubli/intent_classifier"

URDU_AUDIO_MODEL = (
    "JuniKhan123/urdu-deepfake-detector-unified_for_gradio_app"
)
ENGLISH_AUDIO_MODEL = "mo-thecreator/Deepfake-audio-detection"



@spaces.GPU
def _dummy_gpu_check():
    return {"status": "ok"}


# ============================================================
# 1. IMAGE DETECTION (PRETRAINED CONVNEXT MODEL)
# ============================================================

@lru_cache(maxsize=1)
def _load_image_detector():
    """
    Load the existing pretrained ConvNeXt model once per Space process.
    This does not train or fine-tune the model.
    """
    architecture_path = hf_hub_download(
        repo_id=IMAGE_MODEL_REPO,
        filename=IMAGE_ARCHITECTURE,
    )
    checkpoint_path = hf_hub_download(
        repo_id=IMAGE_MODEL_REPO,
        filename=IMAGE_CHECKPOINT,
    )
    transforms_path = hf_hub_download(
        repo_id=IMAGE_MODEL_REPO,
        filename=IMAGE_TRANSFORMS,
    )

    # Load the architecture file from the model repository.
    spec = importlib.util.spec_from_file_location(
        "sentinelzone_convnext_architecture",
        architecture_path,
    )
    if spec is None or spec.loader is None:
        raise RuntimeError("Could not load the ConvNeXt architecture file.")

    architecture_module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = architecture_module
    spec.loader.exec_module(architecture_module)

    build_model = getattr(architecture_module, "build_model", None)
    if build_model is None:
        raise RuntimeError(
            "The downloaded convnext.py does not define build_model()."
        )

    model = build_model()
    checkpoint = torch.load(
        checkpoint_path,
        map_location=torch.device(_DEVICE),
        weights_only=False,
    )

    if not isinstance(checkpoint, dict) or "model" not in checkpoint:
        raise RuntimeError(
            "Unexpected checkpoint format: expected a dictionary with a 'model' key."
        )

    model.load_state_dict(checkpoint["model"])
    model.to(torch.device(_DEVICE))
    model.eval()

    # Import the author's own test transforms so preprocessing stays exact.
    transform_spec = importlib.util.spec_from_file_location(
        "sentinelzone_convnext_transforms",
        transforms_path,
    )
    if transform_spec is None or transform_spec.loader is None:
        raise RuntimeError("Could not load the ConvNeXt preprocessing file.")

    transform_module = importlib.util.module_from_spec(transform_spec)
    sys.modules[transform_spec.name] = transform_module
    transform_spec.loader.exec_module(transform_module)

    test_transforms = getattr(transform_module, "test_transforms", None)
    if test_transforms is None:
        raise RuntimeError(
            "The downloaded transforms.py does not define test_transforms()."
        )

    transform = test_transforms()
    return model, transform


def detect_image(image):
    if image is None:
        return {
            "verdict": "No image",
            "confidence": 0.0,
        }

    try:
        model, transform = _load_image_detector()
        image = image.convert("RGB")
        image_tensor = transform(image).unsqueeze(0).to(_DEVICE)

        with torch.inference_mode():
            logits = model(image_tensor)
            probabilities = torch.softmax(logits.float(), dim=1)[0]

        # Model card/dataset mapping: class 0 = real, class 1 = fake.
        real_probability = float(probabilities[0].item())
        fake_probability = float(probabilities[1].item())

        verdict = (
            "AI-Generated"
            if fake_probability >= 0.50
            else "Real"
        )
        confidence = (
            fake_probability if verdict == "AI-Generated"
            else real_probability
        ) * 100

        return {
            "verdict": verdict,
            "confidence": round(confidence, 2),
            "ai_probability": round(fake_probability * 100, 2),
            "real_probability": round(real_probability * 100, 2),
            "model": IMAGE_MODEL_REPO,
            "predicted_label": "fake" if verdict == "AI-Generated" else "real",
        }

    except Exception as e:
        print(f"ConvNeXt image detection error: {type(e).__name__}: {e}")
        return {
            "verdict": "Error",
            "confidence": 0.0,
            "error": f"{type(e).__name__}: {str(e)}",
            "model": IMAGE_MODEL_REPO,
        }


# ============================================================
# 2. TEXT HARASSMENT DETECTION
# ============================================================

@lru_cache(maxsize=1)
def _load_text_classifier():
    return pipeline(
        "text-classification",
        model=TEXT_MODEL_ID,
        device=-1,
    )


def detect_text(text):
    if not text or not text.strip():
        return {
            "label": "Safe",
            "score": 0.0,
            "risk": "Low",
        }

    text_lower = text.lower().strip()

    SAFE_WORDS = [
        "hello", "hi", "hey", "ok", "okay",
        "thanks", "thank you", "friend", "bhai",
        "bro", "sister", "dost", "kya haal hai",
        "kasy ho", "kasi ho",
    ]

    # Exact phrase matching avoids some substring false positives.
    if (
        len(text_lower.split()) <= 4
        and text_lower in SAFE_WORDS
    ):
        return {
            "label": "Safe",
            "score": 0.0,
            "risk": "Low",
        }

    BAD_WORDS = [
        "harami", "haramkhor", "gandi", "gndi",
        "gaali", "kamini", "kutta", "kutty", "bc",
        "mc", "madarchod", "bhenchod", "behanchod",
        "randi", "raand", "sex", "sexy", "mazy",
        "nudes", "nangi", "besharam", "tharki",
        "lund", "phudi", "chut", "gashti", "kanjar",
        "kanjri", "jahil", "badtameez", "pagal",
        "zaleel", "mot",
    ]

    THREAT_WORDS = [
        "maar dunga", "jaan se", "kill", "murder",
        "tabah", "barbaad", "dekh lunga",
        "chhor dunga", "chhod dunga", "destroy",
        "jaan", "shadi ni krny",
        "mnh dehkany k qabal ni",
    ]

    has_bad = any(
        word in text_lower for word in BAD_WORDS
    )

    has_threat = any(
        word in text_lower for word in THREAT_WORDS
    )

    if has_bad or has_threat:
        return {
            "label": "Harassment",
            "score": 90.0,
            "risk": "High",
        }

    try:
        classifier = _load_text_classifier()
        result = classifier(text)[0]

        label = str(result["label"])
        score_percent = float(result["score"]) * 100

        # Confirm the model's label mapping before production use.
        is_harassment = label.upper() == "LABEL_1"

        if is_harassment:
            if score_percent >= 85:
                risk = "High"
            elif score_percent >= 50:
                risk = "Medium"
            else:
                risk = "Low"

            return {
                "label": (
                    "Harassment"
                    if score_percent >= 50
                    else "Safe"
                ),
                "score": round(score_percent, 2),
                "risk": risk,
            }

        return {
            "label": "Safe",
            "score": round(score_percent, 2),
            "risk": "Low",
        }

    except Exception as e:
        print(f"Text detection error: {e}")

        return {
            "label": "Error",
            "score": 0.0,
            "risk": "Unknown",
            "error": str(e),
        }


# ============================================================
# 3. AUDIO DEEPFAKE DETECTION (URDU & ENGLISH)
# ============================================================

@lru_cache(maxsize=2)
def _load_audio_model(model_id):
    """
    Cache the feature extractor and model.
    Check for uninitialized meta parameters before inference.
    """

    feature_extractor = AutoFeatureExtractor.from_pretrained(
        model_id
    )

    model = AutoModelForAudioClassification.from_pretrained(
        model_id,
        low_cpu_mem_usage=False,
        device_map=None,
        torch_dtype=torch.float32,
    )

    meta_params = [
        name
        for name, param in model.named_parameters()
        if param.is_meta
    ]

    meta_buffers = [
        name
        for name, buffer in model.named_buffers()
        if buffer.is_meta
    ]

    if meta_params or meta_buffers:
        examples = (meta_params + meta_buffers)[:8]

        raise RuntimeError(
            "Audio model has uninitialized meta tensors. "
            f"Examples: {examples}. "
            "Check Transformers, Accelerate, PyTorch, and "
            "checkpoint compatibility in requirements.txt."
        )

    model = model.to(torch.device("cpu"))
    model.eval()

    return feature_extractor, model


def _is_fake_label(label):
    label = str(label).strip().lower().replace("_", "-")

    return any(
        word in label
        for word in (
            "fake",
            "deepfake",
            "synthetic",
            "ai-generated",
            "ai generated",
        )
    )


def _is_real_label(label):
    label = str(label).strip().lower().replace("_", "-")

    return any(
        word in label
        for word in (
            "real",
            "genuine",
            "bonafide",
            "bona-fide",
            "human",
        )
    )


def detect_audio(audio_path, language="Urdu"):
    if not audio_path:
        return {
            "verdict": "No audio",
            "confidence": 0.0,
        }

    model_id = (
        URDU_AUDIO_MODEL
        if language == "Urdu"
        else ENGLISH_AUDIO_MODEL
    )

    try:
        import soundfile as sf

        # Read and convert audio to mono.
        try:
            waveform, sr = sf.read(
                audio_path,
                dtype="float32",
            )

            if waveform.ndim > 1:
                waveform = np.mean(
                    waveform,
                    axis=1,
                    dtype=np.float32,
                )

            if waveform.size == 0:
                return {
                    "verdict": "Error",
                    "confidence": 0.0,
                    "error": "The uploaded audio file is empty.",
                }

            if sr != 16000:
                waveform = librosa.resample(
                    waveform,
                    orig_sr=sr,
                    target_sr=16000,
                )

        except Exception as read_error:
            print(
                "SoundFile read failed; trying librosa:",
                read_error,
            )

            waveform, sr = librosa.load(
                audio_path,
                sr=16000,
                mono=True,
            )

        waveform = np.asarray(
            waveform,
            dtype=np.float32,
        ).reshape(-1)

        if waveform.size == 0:
            return {
                "verdict": "Error",
                "confidence": 0.0,
                "error": (
                    "Audio is empty or could not be decoded."
                ),
            }

        feature_extractor, model = _load_audio_model(
            model_id
        )

        inputs = feature_extractor(
            waveform,
            sampling_rate=16000,
            return_tensors="pt",
            padding=True,
        )

        inputs = {
            key: value.to("cpu")
            for key, value in inputs.items()
        }

        with torch.inference_mode():
            outputs = model(**inputs)
            logits = outputs.logits
            probs = torch.softmax(
                logits.float(),
                dim=-1,
            )[0]

        pred_idx = int(torch.argmax(probs).item())
        confidence = float(
            probs[pred_idx].item() * 100
        )

        id2label = (
            getattr(model.config, "id2label", {}) or {}
        )

        label = id2label.get(
            pred_idx,
            id2label.get(
                str(pred_idx),
                f"LABEL_{pred_idx}",
            ),
        )

        if _is_fake_label(label):
            verdict = "Fake (AI-Generated)"

        elif _is_real_label(label):
            verdict = "Real"

        else:
            # Fallback only for generic labels.
            # Verify the Urdu checkpoint's class mapping
            # against its model card/configuration.
            if language == "English":
                verdict = (
                    "Fake (AI-Generated)"
                    if pred_idx == 0
                    else "Real"
                )
            else:
                verdict = (
                    "Fake (AI-Generated)"
                    if pred_idx == 1
                    else "Real"
                )

        return {
            "verdict": verdict,
            "confidence": round(confidence, 2),
            "model": model_id,
            "predicted_label": str(label),
        }

    except Exception as e:
        print(
            f"Audio detection error ({model_id}): "
            f"{type(e).__name__}: {e}"
        )

        return {
            "verdict": "Error",
            "confidence": 0.0,
            "error": f"{type(e).__name__}: {str(e)}",
            "model": model_id,
        }


# ============================================================
# 4. GRADIO INTERFACE
# ============================================================

with gr.Blocks(title="SentinelZone AI") as demo:

    gr.Markdown(
        "# SentinelZone - AI Detection API\n"
        "Image, text harassment, and audio deepfake detection."
    )

    # IMAGE TAB
    with gr.Tab("Image Detection"):

        img_input = gr.Image(
            type="pil",
            label="Upload Image",
        )

        img_output = gr.JSON(
            label="Image Analysis Result"
        )

        img_btn = gr.Button("Analyze Image")

        img_btn.click(
            fn=detect_image,
            inputs=img_input,
            outputs=img_output,
            api_name="predict_image",
        )

    # TEXT TAB
    with gr.Tab("Text Harassment"):

        txt_input = gr.Textbox(
            label="Enter Text",
            placeholder="Enter the message to analyze...",
            lines=4,
        )

        txt_output = gr.JSON(
            label="Text Analysis Result"
        )

        txt_btn = gr.Button("Analyze Text")

        txt_btn.click(
            fn=detect_text,
            inputs=txt_input,
            outputs=txt_output,
            api_name="predict_text",
        )

    # AUDIO TAB
    with gr.Tab("Audio Deepfake"):

        aud_input = gr.Audio(
            type="filepath",
            label="Upload Audio File",
        )

        aud_lang = gr.Radio(
            choices=["Urdu", "English"],
            label="Select Language",
            value="Urdu",
        )

        aud_output = gr.JSON(
            label="Audio Analysis Result"
        )

        aud_btn = gr.Button("Analyze Audio")

        aud_btn.click(
            fn=detect_audio,
            inputs=[aud_input, aud_lang],
            outputs=aud_output,
            api_name="predict_audio",
        )


if __name__ == "__main__":
    demo.launch()