import logging
import os
import subprocess
import tempfile
import sys
from multiprocessing import freeze_support

import librosa
import numpy as np
import tensorflow_hub as hub
from matplotlib import pyplot as plt
from sklearn.manifold import TSNE
from transformers import pipeline
from voiceauthCore.utils import convert_to_wav, get_file_metadata
from voiceauthCore.database import save_metadata, init_db
from transformers import AutoImageProcessor, AutoModelForImageClassification
from io import BytesIO
from PIL import Image
import requests
import torch
import tensorflow.compat.v1 as tf

freeze_support()
# Load ML models
yamnet_model = hub.load('https://www.kaggle.com/models/google/yamnet/TensorFlow2/yamnet/1')
vggish_model = hub.load("https://www.kaggle.com/models/google/vggish/TensorFlow2/vggish/1")
pipe = pipeline("audio-classification", model="alexandreacff/wav2vec2-large-ft-fake-detection")
pipe2 = pipeline("audio-classification", model="WpythonW/ast-fakeaudio-detector")
pipe3 = pipeline("audio-classification", model="alexandreacff/sew-ft-fake-detection")

# Determine base path (handles both PyInstaller frozen & normal script execution)
if getattr(sys, "frozen", False):
    base_path = os.path.join(tempfile.gettempdir(), "voiceauthCore")  # Temp directory for frozen app
else:
    base_path = os.path.join(os.getcwd(), "voiceauthCore")  # Local directory for regular execution

# Create temp directory once (if it doesn't exist)
os.makedirs(base_path, exist_ok=True)

# Set temp_dir and ensure it's not repeatedly deleted
temp_dir = base_path
ffmpeg_path = os.path.join(base_path, "ffmpeg")

if os.path.exists(ffmpeg_path):
    os.environ["PATH"] += os.pathsep + ffmpeg_path + "ffmpeg.exe"

# Set Librosa cache directory
librosa_cache_dir = os.path.join(tempfile.gettempdir(), "librosa")
os.makedirs(librosa_cache_dir, exist_ok=True)  # Ensure it exists
os.environ["LIBROSA_CACHE_DIR"] = librosa_cache_dir


def setup_logging(log_filename: str = "audio_detection.log") -> None:
    """Sets up logging to both file and console."""
    logging.basicConfig(
        level=logging.DEBUG,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[
            logging.FileHandler(log_filename, mode="a"),
            logging.StreamHandler(),
        ],
    )


setup_logging()  # Call it early in the script
logging.info("App starting...")

# Load the pre-trained model and processor
processor = AutoImageProcessor.from_pretrained("dima806/deepfake_vs_real_image_detection")
model = AutoModelForImageClassification.from_pretrained("dima806/deepfake_vs_real_image_detection")


def analyze_image(image_url):
    response = requests.get(image_url)
    image_bytes = response.content

    image = Image.open(BytesIO(image_bytes)).convert("RGB")
    with torch.no_grad():
        inputs = processor(images=image, return_tensors="pt")
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class = torch.argmax(logits, dim=1).item()
        return predicted_class == 0  # Return True if predicted as real, False otherwise


def analyze_audio(file_path):
    """Analyzes an audio file using all three models and aggregates results."""
    try:
        wav_path = convert_to_wav(file_path)
        label1, confidence1 = predict_yamnet(wav_path)
        label2, confidence2 = predict_vggish(wav_path)
        label3, confidence3 = predict_hf(wav_path)
        label4, confidence4 = predict_hf2(wav_path)
        label5, confidence5 = predict_rf(wav_path)

        results = {
            "YAMNet": {"label": label1, "confidence": confidence1},
            "VGGish": {"label": label2, "confidence": confidence2},
            "HuggingFace Model 1": {"label": label3, "confidence": confidence3},
            "HuggingFace Model 2": {"label": label4, "confidence": confidence4},
            "R-Forest": {"label": label5, "confidence": confidence5},
        }

        save_metadata(file_path, results, max(confidence1, confidence2, confidence3, confidence4, confidence5))
        return results
    except Exception as e:
        return {"error": str(e)}


def class_names_from_csv(class_map_csv_text):
    """Parses class names from YAMNet class map CSV."""
    return [line.split(',')[0] for line in class_map_csv_text.splitlines() if line]

def predict_yamnet(file_path):
    try:
        # Load audio file
        audio, sr = librosa.load(file_path, sr=16000, mono=True)

        # Run model inference
        outputs = yamnet_model(audio)

        # Extract scores, embeddings, and spectrogram
        scores, embeddings, spectrogram = outputs

        # Convert scores to numpy
        scores_np = scores.numpy()

        if scores_np.size == 0:
            raise ValueError("YAMNet model returned empty scores.")

        # Get inferred class (most probable one)
        inferred_class_idx = np.mean(scores_np, axis=0).argmax()

        # Load class names
        class_map_csv_bytes = tf.io.read_file(yamnet_model.class_map_path())
        class_map_text = class_map_csv_bytes.numpy().decode('utf-8')

        # Parse class names
        class_names = class_names_from_csv(class_map_text)

        if inferred_class_idx >= len(class_names):
            raise IndexError("Inferred class index is out of range.")

        # Get class name
        inferred_class_name = class_names[inferred_class_idx]

        return inferred_class_idx, inferred_class_name

    except Exception as e:
        print(f"Error in predict_yamnet: {e}")
        return None, "Unknown"


def predict_vggish(file_path):
    audio, sr = librosa.load(file_path, sr=16000, mono=True)
    embeddings = vggish_model(audio)
    return "VGGish_Features", embeddings.numpy().mean()


def predict_hf(file_path):
    audio_data, sr = librosa.load(file_path, sr=16000)
    prediction = pipe(audio_data)
    return prediction[0]["label"], prediction[0]["score"]


def predict_hf2(file_path):
    audio_data, sr = librosa.load(file_path, sr=16000)
    prediction = pipe2(audio_data)
    return prediction[0]["label"], prediction[0]["score"]


def predict_rf(file_path):
    audio_data, sr = librosa.load(file_path, sr=16000)
    prediction = pipe3(audio_data)
    return prediction[0]["label"], prediction[0]["score"]


def visualize_embeddings_tsne(file_path, output_path="tsne_visualization.png"):
    embeddings = predict_vggish(file_path)

    n_samples = embeddings.shape[0]

    if n_samples <= 1:
        print(
            f"t-SNE cannot be performed with only {n_samples} sample(s). Skipping visualization."
        )

        plt.figure(figsize=(10, 6))
        plt.text(
            0.5,
            0.5,
            "Not enough samples for t-SNE",
            fontsize=12,
            ha="center")
        plt.title("t-SNE Visualization of Audio Embeddings")
        plt.savefig(output_path)
        plt.close()
        os.startfile(output_path)
        return

    perplexity = min(30, n_samples - 1)

    perplexity = max(5.0, perplexity)

    tsne = TSNE(n_components=2, random_state=42, perplexity=perplexity)
    reduced_embeddings = tsne.fit_transform(embeddings)

    plt.figure(figsize=(10, 6))
    plt.scatter(
        reduced_embeddings[:, 0],
        reduced_embeddings[:, 1],
        c="blue",
        alpha=0.7,
        edgecolors="k",
    )
    plt.title("t-SNE Visualization of Audio Embeddings")
    plt.xlabel("Component 1")
    plt.ylabel("Component 2")
    plt.tight_layout()

    plt.savefig(output_path)
    plt.close()

    if sys.platform.system() == "Windows":
        os.startfile(output_path)
    elif sys.platform.system() == "Darwin":  # macOS
        subprocess.run(["open", output_path], check=True)
    else:  # Linux/Unix
        subprocess.run(["xdg-open", output_path], check=True)


if __name__ == "__main__":
    import sys
    file_path = sys.argv[1]
    print(analyze_audio(file_path))

