import datetime
import ntpath
import os
import wave
import tempfile
import shutil
import subprocess
import traceback
import copy
from pyannote.core import Segment

import numpy as np
import torch
from sklearn.cluster import KMeans

from Orange.data import Domain, Table, StringVariable


# --- Fonctions utilitaires ---

def convert_audio_to_pcm(file_path, ffmpeg_path):
    """Convertit un fichier audio vers un format PCM WAV mono 16kHz."""
    print(f"[FFMPEG] Tentative de conversion de: {file_path}")
    try:
        ext = os.path.splitext(file_path)[1].lower()
        if ext not in [".mp3", ".wav", ".m4a"]:
            print(f"[ERREUR] Type de fichier non supporté : {ext}")
            return None

        base_name = os.path.splitext(ntpath.basename(file_path))[0]
        pcm_wav_path = os.path.join(tempfile.gettempdir(), base_name + "_pcm.wav")

        ffmpeg_cmd = [
            ffmpeg_path, "-y",
            "-i", file_path,
            "-acodec", "pcm_s16le",
            "-ar", "16000",
            "-ac", "1",
            pcm_wav_path
        ]

        print(f"[FFMPEG] Commande: {' '.join(ffmpeg_cmd)}")
        result = subprocess.run(ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

        if result.returncode != 0:
            print(f"[ERREUR] ffmpeg a échoué avec le code {result.returncode}")
            print("[STDERR]", result.stderr)
            return None

        if os.path.exists(pcm_wav_path):
            print(f"[FFMPEG] Conversion réussie vers: {pcm_wav_path}")
            return pcm_wav_path
        else:
            print("[ERREUR] Le fichier converti n’a pas été trouvé.")
            return None

    except Exception as e:
        print(f"❌ Exception pendant la conversion audio : {e}")
        return None


def get_wav_duration(wav_path):
    """Calcule la durée d'un fichier WAV."""
    try:
        with wave.open(wav_path, "rb") as wav_file:
            frames = wav_file.getnframes()
            rate = wav_file.getframerate()
            duration = frames / float(rate)
            print(f"[INFO] Durée du WAV calculée : {duration:.2f} secondes")
            return duration
    except Exception as e:
        print(f"An error occurred when getting the wav duration : {e}")
        return None


def process_input_file(file_path, ffmpeg_path):
    """Copie le fichier d'entrée vers un emplacement temporaire et le convertit en PCM."""
    print(f"[PREP] Démarrage du processus de préparation pour: {file_path}")
    try:
        temp_dir = tempfile.gettempdir()
        base_name = "input_audio_" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        base_ext = os.path.splitext(file_path)[1]
        clean_copy = os.path.join(temp_dir, base_name + base_ext)

        try:
            shutil.copy(file_path, clean_copy)
            file_to_convert = clean_copy
            print(f"[PREP] Fichier copié temporairement: {clean_copy}")
        except Exception as e:
            file_to_convert = file_path
            print(f"[PREP] Avertissement: Échec de la copie ({e}). Utilisation du chemin original.")

        pcm_path = convert_audio_to_pcm(file_to_convert, ffmpeg_path)

        if pcm_path and file_to_convert != file_path:
            # Nettoyage de la copie temporaire si elle a été faite
            os.remove(file_to_convert)
            print(f"[PREP] Nettoyage de la copie temporaire: {file_to_convert}")

        return pcm_path
    except Exception as e:
        print(f"[ERREUR] Erreur pendant le traitement du fichier: {e}")
        return None


def transcription(table, model, embedding_model, audio_helper, ffmpeg_path, column_name, progress_callback=None,
                  num_speakers=2, argself=None):
    """
    Exécute la transcription et la diarisation de manière synchrone.
    NOTE: La logique actuelle itère sur la table mais utilise seulement le dernier file_path pour la transcription.
    """
    if table is None or len(table) == 0:
        print("[PROCESS] Table d'entrée vide. Retourne None.")
        return None


    data = copy.deepcopy(table)
    attr_dom = list(data.domain.attributes)
    metas_dom = list(data.domain.metas)
    class_dom = list(data.domain.class_vars)

    file_path = None

    print(f"[PROCESS] Démarrage de la boucle sur {len(data)} lignes.")
    for i, row in enumerate(data):
        print(f"[ROW {i + 1}/{len(data)}] Début du traitement de la ligne.")

        _file_path = str(row[column_name])
        print(f"[ROW {i + 1}] Fichier audio source: {_file_path}")

        # Préparation du fichier PCM
        file_path = process_input_file(_file_path, ffmpeg_path)

        if file_path is None:
            print(
                f"[ROW {i + 1}] Avertissement: Échec de la préparation du fichier. Passe à la ligne suivante (ou arrête si une seule ligne).")
            # Si c'était un vrai traitement batch, on continuerait ici.
            # Mais comme la fonction ne retourne qu'un seul résultat, nous allons nous en tenir au dernier path valide.

        if progress_callback is not None:
            progress_value = float(100 * (i + 1) / len(data))
            progress_callback(progress_value)

        if argself is not None:
            if argself.stop:
                print(f"[PROCESS] Arrêt demandé après la ligne {i + 1}.")
                break

        print(f"[ROW {i + 1}] Fin de la préparation de la ligne.")
    if file_path is None:
        return "Erreur : Aucun fichier audio n'a pu être préparé correctement.", [], 0.0

    print(f"\n[PROCESS] Démarrage de la transcription pour le fichier PCM final: {file_path}")

    try:
        if not os.path.exists(file_path):
            print(f"[ERREUR] Fichier PCM final introuvable: {file_path}")
            return "Erreur : fichier introuvable.", [], 0.0

        file_duration = get_wav_duration(file_path)
        if file_duration is None:
            return "Error: File duration unknown.", [], 0.0

        start_time = datetime.datetime.now()

        # 1. Transcription (Whisper)
        print("[WHISPER] Démarrage de la transcription...")
        result = model.transcribe(
            file_path, language="fr", without_timestamps=False, temperature=0
        )

        print(f"[WHISPER] Transcription terminée. Segments trouvés: {len(result.get('segments', []))}")

        transcription_time = (datetime.datetime.now() - start_time).total_seconds() / 60

        if not result or "segments" not in result or not result["segments"]:
            print("[WHISPER] ERREUR: Aucun segment de parole détecté.")
            return "Error: No speech detected.", [], transcription_time

        segments = result["segments"]
        embeddings = None
        valid_segments_indices = []

        # 2. Extraction des Embeddings (Diarisation)
        print("[DIARIZATION] Démarrage de l'extraction des embeddings...")
        for i, segment in enumerate(segments):
            start, end = segment["start"], segment["end"]

            if end > file_duration:
                print(f"[EMBED] Segment {i} ignoré (fin {end:.2f} > durée {file_duration:.2f})")
                continue

            valid_segments_indices.append(i)
            try:
                # Découpage du segment
                waveform, _ = audio_helper.crop(file_path, Segment(start, end))
                if waveform.ndim == 1: waveform = waveform.unsqueeze(0)

                # Encodage
                with torch.no_grad():
                    embedding = embedding_model.encode_batch(waveform).squeeze().cpu().numpy()

                if embeddings is None:
                    embeddings = np.zeros((len(segments), embedding.shape[0]))
                    print(f"[EMBED] Initialisation de l'array d'embeddings (Dim: {embedding.shape[0]})")

                embeddings[i] = embedding

            except Exception as e:
                print(f"[EMBED] ERREUR lors du traitement du segment {i} ({start:.2f}-{end:.2f}s): {e}")
                if i in valid_segments_indices: valid_segments_indices.remove(i)
                continue

        num_valid_segments = len(valid_segments_indices)
        print(f"[DIARIZATION] Extraction terminée. Segments valides pour le clustering: {num_valid_segments}")

        if embeddings is None or embeddings.size == 0 or num_valid_segments == 0:
            return "Error: Could not extract valid embeddings.", [], transcription_time

        valid_segments = [segments[i] for i in valid_segments_indices]
        valid_embeddings = embeddings[valid_segments_indices]

        # 3. Clustering
        n_clusters = min(num_speakers, num_valid_segments)
        if n_clusters < 1:
            print("[CLUSTER] ERREUR: Moins d'un segment valide pour le clustering.")
            return "Error: Not enough segments for clustering.", [], transcription_time

        print(f"[CLUSTER] Démarrage de KMeans avec {n_clusters} locuteur(s) ciblé(s).")
        clustering = KMeans(n_clusters=n_clusters, random_state=42, n_init='auto').fit(valid_embeddings)
        labels = clustering.labels_
        print("[CLUSTER] Clustering terminé.")

        # 4. Regroupement des segments
        speaker_map = {}
        merged_segments = []
        current_speaker = None
        current_text = ""
        current_start = None
        table_output = []

        print("[MERGE] Démarrage du regroupement des segments par locuteur.")
        for i, segment in enumerate(valid_segments):
            speaker_id = labels[i]
            if speaker_id not in speaker_map:
                speaker_map[speaker_id] = f"SPEAKER {len(speaker_map) + 1}"
            speaker = speaker_map[speaker_id]

            # Détection de changement de locuteur ou début
            if current_speaker != speaker:
                if current_speaker is not None:
                    # Finaliser le bloc précédent
                    timestamp = str(datetime.timedelta(seconds=round(current_start)))
                    merged_segments.append(f"{current_speaker} {timestamp}: {current_text}")
                    table_output.append([current_speaker, timestamp, current_text])

                # Démarrer le nouveau bloc
                current_speaker = speaker
                current_text = segment["text"]
                current_start = segment["start"]

            else:
                # Continuer le bloc actuel
                current_text += f" {segment['text']}"

        # Finaliser le dernier bloc
        if current_speaker is not None:
            timestamp = str(datetime.timedelta(seconds=round(current_start)))
            merged_segments.append(f"{current_speaker} {timestamp}: {current_text}")
            table_output.append([current_speaker, timestamp, current_text])

        answer_dom = [StringVariable("Answer")]
        timestamp_dom = [StringVariable("timestamp")]
        current_text_dom = [StringVariable("current_text")]
        domain = Domain(attributes=attr_dom , class_vars=class_dom, metas=metas_dom + answer_dom + timestamp_dom + current_text_dom)

        # Create and return table
        out_data = Table.from_list(domain=domain, rows=table_output)
        return out_data

    except Exception as e:
        print(f"[ERREUR FATALE] Exception non gérée durant la transcription/diarisation: {e}")
        traceback.print_exc()
        return f"Erreur fatale: {e}", [], 0.0