"""Dataset exporter for converting Revoxx sessions to Talrómur 3 format."""

import shutil
import json
from pathlib import Path
from typing import List, Dict, Set, Tuple, Optional, Any
from collections import Counter
import soundfile as sf

from ..session.inspector import SessionInspector
from ..session.script_parser import FestivalScriptParser
from ..controllers.session_controller import REFERENCE_SILENCE_LABEL


class DatasetExporter:
    """Export Revoxx sessions to Talrómur 3 dataset format.

    The Talrómur 3 format organizes recordings by emotion with an index.tsv file:
    - voice-name/emotion/voice-name_emotion_001.wav
    - voice-name/index.tsv with metadata

    Intensity Level Handling:
    - Text format "N: text" where N (1-5) is the intensity level
    - Neutral emotion always gets intensity 0 (regardless of text prefix)
    - Other emotions preserve the intensity from the text prefix
    - Missing intensity prefix defaults to 0
    """

    def __init__(
        self,
        output_dir: Path,
        audio_format: str = "flac",
        zero_intensity_emotions: List[str] = None,
        include_intensity: bool = True,
        include_vad: bool = False,
    ):
        """Initialize dataset exporter.

        Args:
            output_dir: Base output directory for datasets
            audio_format: Output audio format ('wav' or 'flac')
            zero_intensity_emotions: List of emotions to set intensity to 0
            include_intensity: Whether to include intensity column in index.tsv
            include_vad: Whether to run VAD analysis on the exported dataset
        """
        self.output_dir = Path(output_dir)
        self.format = audio_format.lower()
        self.zero_intensity_emotions = zero_intensity_emotions or ["neutral"]
        self.include_intensity = include_intensity
        self.include_vad = include_vad

    def _group_sessions_by_speaker(self, session_paths: List[Path]) -> Dict:
        """Group sessions by speaker name.

        Args:
            session_paths: List of session paths

        Returns:
            Dictionary mapping speaker names to session data
        """
        speaker_groups = {}
        for session_path in session_paths:
            session_data = self._load_session(session_path)
            if session_data:
                speaker_name = session_data.get("speaker", {}).get("name", "unknown")
                speaker_name_normalized = speaker_name.lower().replace(" ", "_")

                if speaker_name_normalized not in speaker_groups:
                    speaker_groups[speaker_name_normalized] = []
                speaker_groups[speaker_name_normalized].append(
                    (session_path, session_data)
                )
        return speaker_groups

    def export_sessions(
        self,
        session_paths: List[Path],
        dataset_name: str = None,
        progress_callback=None,
        skip_rejected: bool = False,
    ) -> Tuple[List[Path], Dict]:
        """Export Revoxx sessions grouped by speaker name.

        Sessions with the same speaker name will be grouped into one dataset.

        Args:
            session_paths: List of paths to .revoxx session directories
            dataset_name: Optional override for dataset name (if None, uses speaker names)
            progress_callback: Optional callback for progress updates
            skip_rejected: If True, skip utterances flagged as 'rejected'

        Returns:
            Tuple of (list of output_paths, statistics_dict)
        """
        if not session_paths:
            raise ValueError("No sessions provided")

        # Load session metadata and group by speaker name
        speaker_groups = self._group_sessions_by_speaker(session_paths)

        if not speaker_groups:
            raise ValueError("No valid sessions found")

        # Process each speaker group as a separate dataset
        all_datasets = []
        total_statistics = {
            "total_utterances": 0,
            "missing_recordings": 0,
            "datasets_created": 0,
            "sessions_processed": 0,
            "speakers": [],
        }

        for speaker_name_normalized, sessions in speaker_groups.items():
            # Use override name if provided and only one speaker group
            if dataset_name and len(speaker_groups) == 1:
                current_dataset_name = dataset_name
            else:
                current_dataset_name = speaker_name_normalized

            # Create output directory for this speaker
            dataset_dir = self.output_dir / current_dataset_name
            if dataset_dir.exists():
                # For now, we'll overwrite. In UI, we'll ask for confirmation
                shutil.rmtree(dataset_dir)
            dataset_dir.mkdir(parents=True, exist_ok=True)

            # Process sessions for this speaker
            index_data = []
            file_counts = Counter()
            speaker_statistics = {
                "speaker": speaker_name_normalized,
                "emotions": set(),
                "sessions": len(sessions),
            }

            # Collect rejected labels from all sessions if skip_rejected is enabled
            rejected_labels = None
            if skip_rejected:
                rejected_labels = set()
                for session_path, session_data in sessions:
                    flags = session_data.get("utterance_flags", {})
                    for label, flag_type in flags.items():
                        if flag_type == "rejected":
                            rejected_labels.add(label)

            # Group sessions by emotion
            emotion_sessions = {}
            for session_path, session_data in sessions:
                emotion = session_data.get("speaker", {}).get("emotion", "unknown")
                if emotion not in emotion_sessions:
                    emotion_sessions[emotion] = []
                emotion_sessions[emotion].append((session_path, session_data))
                speaker_statistics["emotions"].add(emotion)

            # Process each emotion group
            for emotion, emotion_session_list in emotion_sessions.items():
                self._process_emotion_group(
                    emotion,
                    emotion_session_list,
                    dataset_dir,
                    current_dataset_name,
                    index_data,
                    file_counts,
                    total_statistics,
                    progress_callback,
                    rejected_labels=rejected_labels,
                )

            # Write index file for this speaker
            index_path = dataset_dir / "index.tsv"
            with open(index_path, "w", encoding="utf-8") as f:
                f.writelines(index_data)

            # Get audio properties from first session
            audio_properties = self._get_audio_properties(sessions[0][1])

            # Write README file with format documentation
            self._write_readme(dataset_dir, audio_properties)

            all_datasets.append(dataset_dir)
            total_statistics["datasets_created"] += 1
            total_statistics["sessions_processed"] += len(sessions)
            total_statistics["speakers"].append(
                {
                    "name": speaker_name_normalized,
                    "emotions": list(speaker_statistics["emotions"]),
                    "file_counts": dict(file_counts),
                    "output_path": str(dataset_dir),
                }
            )

        # Run VAD processing if requested
        if self.include_vad:
            vad_stats = self._run_vad_processing(all_datasets, progress_callback)
            total_statistics["vad_statistics"] = vad_stats

        return all_datasets, total_statistics

    def _process_emotion_group(
        self,
        emotion: str,
        emotion_session_list: List,
        dataset_dir: Path,
        dataset_name: str,
        index_data: List,
        file_counts: Counter,
        total_statistics: Dict,
        progress_callback=None,
        rejected_labels: Optional[Set[str]] = None,
    ) -> None:
        """Process all sessions for a specific emotion.

        Args:
            emotion: Emotion name
            emotion_session_list: List of sessions for this emotion
            dataset_dir: Output directory for dataset
            dataset_name: Name of the dataset
            index_data: List to append index entries to
            file_counts: Counter for file numbering
            total_statistics: Statistics dictionary to update
            progress_callback: Optional progress callback
            rejected_labels: Set of utterance labels to skip during export
        """
        emotion_dir = dataset_dir / emotion
        emotion_dir.mkdir(exist_ok=True)

        # Process all utterances for this emotion
        utterance_map = self._collect_utterances(
            emotion_session_list, rejected_labels=rejected_labels
        )

        for utterance_id, (session_path, take_num, text) in utterance_map.items():
            total_statistics["total_utterances"] += 1

            # Extract intensity and clean text
            intensity, clean_text = self._extract_intensity_and_text(text)
            if emotion in self.zero_intensity_emotions:
                intensity = "0"

            # Process audio file
            source_file = (
                session_path / "recordings" / utterance_id / f"take_{take_num:03d}.flac"
            )

            if source_file.exists():
                file_counter = file_counts[emotion] + 1
                output_filename = (
                    f"{dataset_name}_{emotion}_{file_counter:03d}.{self.format}"
                )
                output_path = emotion_dir / output_filename

                # Copy or convert audio
                if self.format == "flac" and source_file.suffix == ".flac":
                    shutil.copy2(source_file, output_path)
                else:
                    self._convert_audio(source_file, output_path)

                file_counts[emotion] += 1

                # Add to index
                if self.include_intensity:
                    index_data.append(
                        f"{output_filename}\t{dataset_name}\t{emotion}\t{intensity}\t{clean_text}\n"
                    )
                else:
                    index_data.append(
                        f"{output_filename}\t{dataset_name}\t{emotion}\t{clean_text}\n"
                    )
            else:
                total_statistics["missing_recordings"] += 1

            if progress_callback:
                progress_callback(total_statistics["total_utterances"])

    @staticmethod
    def _load_session(session_path: Path) -> Optional[Dict]:
        """Load session metadata from session.json."""
        return SessionInspector.load_metadata(session_path)

    def _collect_utterances(
        self,
        emotion_sessions: List[Tuple[Path, Dict]],
        rejected_labels: Optional[Set[str]] = None,
    ) -> Dict:
        """Collect all utterances from sessions, choosing best take for each.

        Args:
            emotion_sessions: List of (session_path, session_data) tuples
            rejected_labels: Set of utterance labels to skip during export

        Returns:
            Dict mapping utterance_id to (session_path, take_number, text)
        """
        utterances = {}

        for session_path, session_data in emotion_sessions:
            recordings_dir = session_path / "recordings"
            script_file = session_path / "script.txt"
            script_data = self._load_script(script_file)

            # Find all recordings
            if recordings_dir.exists():
                for utterance_dir in recordings_dir.iterdir():
                    if utterance_dir.is_dir():
                        utterance_id = utterance_dir.name

                        # Skip reference silence - it's not part of the dataset
                        if utterance_id == REFERENCE_SILENCE_LABEL:
                            continue

                        # Skip rejected utterances
                        if rejected_labels and utterance_id in rejected_labels:
                            continue

                        # Find highest take number
                        takes = list(utterance_dir.glob("take_*.flac"))
                        takes.extend(list(utterance_dir.glob("take_*.wav")))

                        if takes:
                            # Extract take numbers and find highest
                            take_numbers = []
                            for take_file in takes:
                                try:
                                    take_num = int(take_file.stem.split("_")[1])
                                    take_numbers.append(take_num)
                                except (IndexError, ValueError):
                                    continue

                            if take_numbers:
                                highest_take = max(take_numbers)
                                text = script_data.get(utterance_id, "")
                                utterances[utterance_id] = (
                                    session_path,
                                    highest_take,
                                    text,
                                )

        return utterances

    @staticmethod
    def _load_script(script_file: Path) -> Dict[str, str]:
        """Load script file and return mapping of utterance_id to text."""
        return FestivalScriptParser.parse_script(script_file)

    @staticmethod
    def _get_audio_properties(session_data: Dict) -> Dict[str, Any]:
        """Extract audio properties from session data.

        Returns:
            Dict with sample_rate and bit_depth
        """
        audio_config = session_data.get("audio_config", {})
        return {
            "sample_rate": audio_config.get("sample_rate"),
            "bit_depth": audio_config.get("bit_depth"),
        }

    @staticmethod
    def _extract_intensity_and_text(text: str) -> Tuple[str, str]:
        """Extract intensity level and clean text from utterance.

        The text format can be:
        - "N: actual text" where N is intensity level 1-5
        - Just "actual text" for no intensity (defaults to "0")

        Note: The calling code may override the intensity to "0" for
        specific emotions like "neutral" regardless of the text prefix.

        Returns:
            Tuple of (intensity, clean_text)
        """
        return FestivalScriptParser.extract_intensity_and_text(text)

    @staticmethod
    def _convert_audio(source_path: Path, dest_path: Path):
        """Convert audio file to target format."""
        try:
            data, samplerate = sf.read(str(source_path))
            sf.write(str(dest_path), data, samplerate)
        except Exception:
            # If conversion fails, try to copy as-is
            shutil.copy2(source_path, dest_path)

    def _write_readme(self, dataset_dir: Path, audio_properties: Dict[str, Any]):
        """Write README.txt with TSV format documentation from template.

        Args:
            dataset_dir: Directory to write README to
            audio_properties: Dict with sample_rate and bit_depth
        """
        # Load main template
        template_dir = Path(__file__).parent.parent / "resources" / "templates"
        main_template_path = template_dir / "dataset_readme.txt"

        with open(main_template_path, "r", encoding="utf-8") as f:
            readme_content = f.read()

        # Load appropriate index format template
        if self.include_intensity:
            index_template_path = template_dir / "index_format_with_intensity.txt"
        else:
            index_template_path = template_dir / "index_format_without_intensity.txt"

        with open(index_template_path, "r", encoding="utf-8") as f:
            index_format = f.read()

        # Prepare audio properties strings
        sample_rate_str = (
            str(audio_properties["sample_rate"])
            if audio_properties["sample_rate"]
            else "Not specified"
        )
        bit_depth_str = (
            str(audio_properties["bit_depth"])
            if audio_properties["bit_depth"]
            else "Not specified"
        )

        # Fill in template variables
        readme_content = readme_content.format(
            index_format=index_format,
            audio_format=self.format.upper(),
            file_extension=self.format,
            sample_rate=sample_rate_str,
            bit_depth=bit_depth_str,
        )

        # Write README file
        readme_path = dataset_dir / "README.txt"
        with open(readme_path, "w", encoding="utf-8") as f:
            f.write(readme_content)

    def _run_vad_processing(
        self, dataset_paths: List[Path], progress_callback=None
    ) -> Dict:
        """Run VAD processing on exported datasets using multiprocessing.

        Args:
            dataset_paths: List of dataset directories to process
            progress_callback: Optional progress callback (count, message)

        Returns:
            Dictionary with total files processed and warnings
        """
        try:
            from scripts_module.vadiate import get_audio_files
            import multiprocessing as mp
            from concurrent.futures import ProcessPoolExecutor, as_completed
        except ImportError:
            return {}  # VAD not available

        # Count total files for progress
        total_files = sum(len(get_audio_files(str(d))) for d in dataset_paths)
        if total_files == 0:
            return {}

        processed = 0
        vad_statistics = {"total_files": total_files, "warnings": []}

        # Use process pool for parallel processing
        # Each process handles VAD analysis for one complete dataset (speaker)
        # This means if we export 3 speakers, we use up to 3 processes
        # Each process analyzes all audio files within its assigned speaker's dataset
        num_workers = min(mp.cpu_count(), len(dataset_paths))

        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            # Submit one VAD processing task per dataset (per speaker)
            # Each task processes all audio files in that speaker's dataset directory
            future_to_dataset = {
                executor.submit(self._process_dataset_vad, dataset_path): dataset_path
                for dataset_path in dataset_paths
            }

            # Process completed tasks
            for future in as_completed(future_to_dataset):
                dataset_path = future_to_dataset[future]
                try:
                    result = future.result()
                    processed += result["files_processed"]
                    vad_statistics["warnings"].extend(result["warnings"])
                    if progress_callback:
                        progress_callback(
                            processed, f"VAD analysis: {processed}/{total_files}"
                        )
                except Exception as e:
                    vad_statistics["warnings"].append(
                        f"VAD processing error for {dataset_path}: {e}"
                    )

        return vad_statistics

    @staticmethod
    def _process_dataset_vad(dataset_path: Path) -> Dict:
        """Process VAD for a single dataset (one speaker's complete dataset).

        This method runs in a separate process and handles all audio files
        for one speaker. If multiple speakers were exported, each speaker's
        dataset is processed by a different process in parallel.

        Args:
            dataset_path: Path to the dataset directory for one speaker

        Returns:
            Dictionary with files processed and warnings
        """
        from scripts_module.vadiate import (
            get_audio_files,
            process_audio,
            load_silero_vad,
        )

        vad_output = dataset_path / "vad.json"
        audio_files = get_audio_files(str(dataset_path))

        result_info = {"files_processed": 0, "warnings": []}

        if not audio_files:
            return result_info

        # Load model for this process
        model = load_silero_vad()
        results = {}

        for file_path in audio_files:
            try:
                rel_path, result, warnings = process_audio(
                    file_path,
                    model,
                    str(dataset_path),
                    use_dynamic_threshold=True,
                    collect_warnings=True,
                )
                results[rel_path] = result
                result_info["warnings"].extend(warnings)
                result_info["files_processed"] += 1
            except Exception as e:
                result_info["warnings"].append(f"VAD error for {file_path}: {e}")

        # Save results
        with open(vad_output, "w") as f:
            json.dump(results, f, indent=2)

        return result_info
