package config

import (
    "os"
    "path/filepath"
)

type Diarization struct {
    Enabled             bool    `yaml:"enabled"`
    ModelDir            string  `yaml:"model_dir"`
    SegmentationModel   string  `yaml:"segmentation_model"`
    EmbeddingModel      string  `yaml:"embedding_model"`
    NumThreads          int     `yaml:"num_threads"`
    NumClusters         int     `yaml:"num_clusters"`
    ClusteringThreshold float32 `yaml:"clustering_threshold"`
    MinDurationOn       float32 `yaml:"min_duration_on"`
    MinDurationOff      float32 `yaml:"min_duration_off"`
}

func (d Diarization) WithDefaults() Diarization {
    if d.ModelDir == "" {
        d.ModelDir = "./models/diarization"
    }
    if d.SegmentationModel == "" {
        d.SegmentationModel = "pyannote-segmentation-3-0/model.onnx"
    }
    if d.EmbeddingModel == "" {
        d.EmbeddingModel = "3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx"
    }
    if d.NumThreads <= 0 {
        d.NumThreads = 2
    }
    if d.ClusteringThreshold <= 0 {
        d.ClusteringThreshold = 0.5
    }
    if d.MinDurationOn <= 0 {
        d.MinDurationOn = 0.3
    }
    if d.MinDurationOff <= 0 {
        d.MinDurationOff = 0.5
    }
    return d
}

func (d Diarization) SegmentationPath() string {
    return resolveModelPath(d.ModelDir, d.SegmentationModel)
}

func (d Diarization) EmbeddingPath() string {
    return resolveModelPath(d.ModelDir, d.EmbeddingModel)
}

func (d Diarization) ModelsPresent() bool {
    if _, err := os.Stat(d.SegmentationPath()); err != nil {
        return false
    }
    if _, err := os.Stat(d.EmbeddingPath()); err != nil {
        return false
    }
    return true
}

func resolveModelPath(dir, name string) string {
    if filepath.IsAbs(name) {
        return name
    }
    return filepath.Join(dir, name)
}

func (d Diarization) Active() bool {
    return d.Enabled
}