"""
Factory functions for creating ONNX inference sessions and models.

This module provides helper functions to create InferenceSessionProxy instances
with appropriate configuration for different model types, as well as convenience
functions to construct complete model instances and pipelines.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from onnx_dl import InferenceSessionProxy
import onnxruntime as ort

from onnx_diarization.clustering.plda import PLDATransform, load_xvex_and_plda_data
from onnx_diarization.embedding import WeSpeakerEmbeddingModel
from onnx_diarization.fbank import FbankExtractor
from onnx_diarization.segmentation import PyannnoteSegmentation

if TYPE_CHECKING:
    from pathlib import Path

    from onnx_diarization.pipeline import SpeakerDiarizationPipeline


def create_onnx_session_options() -> ort.SessionOptions:
    """
    Create ONNX Runtime session options optimized for memory management.

    Returns:
        Session options with memory arena disabled for better memory control
    """
    sess_options = ort.SessionOptions()
    sess_options.enable_cpu_mem_arena = False
    sess_options.enable_mem_pattern = False
    return sess_options


def create_cuda_provider_options() -> dict[str, str | int]:
    """
    Create CUDA provider options with memory limits.

    Returns:
        Dictionary of CUDA-specific provider options
    """
    return {
        "arena_extend_strategy": "kSameAsRequested",
        "cudnn_conv_algo_search": "DEFAULT",
        "do_copy_in_default_stream": "1",
    }


def create_onnx_providers() -> list[tuple[str, dict[str, str | int]] | str]:
    """
    Create ONNX providers with memory-optimized settings.

    Returns:
        List of execution providers, trying CUDA first then falling back to CPU
    """
    return [
        ("CUDAExecutionProvider"),
        # ("CUDAExecutionProvider", create_cuda_provider_options()),
        # "CPUExecutionProvider",
    ]


def create_segmentation_session(model_path: Path | str, ttl: int = -1) -> InferenceSessionProxy | ort.InferenceSession:
    """
    Create InferenceSessionProxy for segmentation model.

    Args:
        model_path: Path to the ONNX segmentation model file
        ttl: Time-to-live in seconds for automatic model unloading
            -1: Never unload (persistent, recommended for production)
             0: Unload immediately after each use (minimal memory)
            >0: Unload after N seconds of inactivity

    Returns:
        InferenceSessionProxy configured for segmentation model
    """
    return ort.InferenceSession(
        model_path, sess_options=create_onnx_session_options(), providers=create_onnx_providers()
    )
    # return InferenceSessionProxy(
    #     model_path, sess_options=create_onnx_session_options(), providers=create_onnx_providers(), ttl=ttl
    # )


def create_embedding_session(model_path: Path | str, ttl: int = -1) -> InferenceSessionProxy | ort.InferenceSession:
    """
    Create InferenceSessionProxy for embedding model.

    Args:
        model_path: Path to the ONNX embedding model file
        ttl: Time-to-live in seconds for automatic model unloading
            -1: Never unload (persistent, recommended for production)
             0: Unload immediately after each use (minimal memory)
            >0: Unload after N seconds of inactivity

    Returns:
        InferenceSessionProxy configured for embedding model
    """
    return ort.InferenceSession(
        model_path, sess_options=create_onnx_session_options(), providers=create_onnx_providers()
    )
    # return InferenceSessionProxy(
    #     model_path, sess_options=create_onnx_session_options(), providers=create_onnx_providers(), ttl=ttl
    # )


def create_segmentation_model(
    model_path: Path | str, duration: float = 10.0, step: float | None = None, ttl: int = -1
) -> PyannnoteSegmentation:
    """
    Create a complete segmentation model with session.

    Args:
        model_path: Path to the ONNX segmentation model file
        duration: Duration of segmentation chunks in seconds
        step: Step size between chunks (defaults to duration * 0.1)
        ttl: Time-to-live for session management

    Returns:
        Configured PyannnoteSegmentation instance
    """
    session = create_segmentation_session(model_path, ttl=ttl)
    return PyannnoteSegmentation(session, duration=duration, step=step)


def create_embedding_model(model_path: Path | str, ttl: int = -1) -> WeSpeakerEmbeddingModel:
    """
    Create a complete embedding model with session.

    Args:
        model_path: Path to the ONNX embedding model file
        ttl: Time-to-live for session management

    Returns:
        Configured WeSpeakerEmbeddingModel instance
    """
    session = create_embedding_session(model_path, ttl=ttl)
    fbank_extractor = FbankExtractor(sample_rate=16000)
    return WeSpeakerEmbeddingModel(session, fbank_extractor)


def create_pipeline(
    segmentation_model_path: Path | str,
    embedding_model_path: Path | str,
    segmentation_duration: float = 10.0,
    segmentation_step: float = 0.1,
    clustering_threshold: float = 0.6,
    embedding_batch_size: int = 128,
    embedding_exclude_overlap: bool = False,
    ttl: int = -1,
) -> SpeakerDiarizationPipeline:
    """
    Factory function to create a complete speaker diarization pipeline.

    This convenience function creates all necessary components (sessions, models,
    and pipeline) with sensible defaults.

    Args:
        segmentation_model_path: Path to segmentation ONNX model
        embedding_model_path: Path to embedding ONNX model
        segmentation_duration: Duration of segmentation chunks in seconds
        segmentation_step: Step size as fraction of duration
        clustering_threshold: Threshold for speaker clustering
        embedding_batch_size: Batch size for embedding extraction
        embedding_exclude_overlap: Whether to exclude overlapping speech
        ttl: Time-to-live for all models (-1=persistent, 0=immediate, >0=seconds)

    Returns:
        Fully configured SpeakerDiarizationPipeline instance
    """
    # Import here to avoid circular dependency
    from onnx_diarization.pipeline import SpeakerDiarizationPipeline

    # Create models
    segmentation = create_segmentation_model(
        segmentation_model_path,
        duration=segmentation_duration,
        step=segmentation_step * segmentation_duration,
        ttl=ttl,
    )
    embedding = create_embedding_model(embedding_model_path, ttl=ttl)

    # Load PLDA and x-vector transform data
    xvec_data, plda_data = load_xvex_and_plda_data()
    plda = PLDATransform(xvec_data, plda_data)

    # Create pipeline
    return SpeakerDiarizationPipeline(
        segmentation=segmentation,
        embedding=embedding,
        plda=plda,
        clustering_threshold=clustering_threshold,
        embedding_batch_size=embedding_batch_size,
        embedding_exclude_overlap=embedding_exclude_overlap,
    )
