package model

import (
	"fmt"
	"sort"
)

const (
	BackendSenseVoice = "sensevoice"
	BackendMoonshine  = "moonshine"
)

var modelProfiles = map[string]ModelProfile{
	DefaultModelID: {
		ID:          DefaultModelID,
		DisplayName: "SenseVoice",
		BackendKind: BackendSenseVoice,
		Tier:        "default",
		SupportedLanguageIDs: []string{
			"zh-CN",
			"zh-Hans",
			"zh-Hant",
			"yue",
			"en",
			"ja",
			"ko",
		},
		RecommendedFor: []string{"zh-CN"},
		Description:    "轻量中文模型，适合中文和中英混输。",
		DownloadURLs: []string{
			"http://www.maikami.com/voicesnap/sensevoice.zip",
			"https://modelscope.cn/models/sherpa-onnx/sherpa-onnx-sense-voice-zh-en-ja-ko-yue/resolve/master/sherpa-onnx-sense-voice-zh-en-ja-ko-yue-int8-2024-07-17.tar.bz2",
		},
		InstallDirName: "sensevoice-zh",
		LegacyDirNames: []string{"sensevoice"},
		RequiredFiles: []RequiredFileRule{
			{
				Role:     "model",
				AnyOf:    []string{"model.int8.onnx", "model.onnx"},
				Required: true,
			},
			{
				Role:     "tokens",
				AllOf:    []string{"tokens.txt"},
				Required: true,
			},
		},
		ProviderOrder: []string{"coreml", "cpu"},
		NumThreads:    4,
	},
	MoonshineModelID: {
		ID:                   MoonshineModelID,
		DisplayName:          "Moonshine English",
		BackendKind:          BackendMoonshine,
		Tier:                 "default",
		SupportedLanguageIDs: []string{EnglishLanguageID},
		RecommendedFor:       []string{EnglishLanguageID},
		Description:          "轻量英文离线模型，适合英文听写和英文日常输入。",
		ApproxSize:           "约 239 MiB 下载包，解压后约 270-280 MiB",
		DownloadURLs: []string{
			"https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-moonshine-base-en-int8.tar.bz2",
		},
		InstallDirName: "moonshine-en",
		RequiredFiles: []RequiredFileRule{
			{
				Role:     "preprocessor",
				AllOf:    []string{"preprocess.onnx"},
				Required: true,
			},
			{
				Role:     "encoder",
				AllOf:    []string{"encode.int8.onnx"},
				Required: true,
			},
			{
				Role:     "uncached_decoder",
				AllOf:    []string{"uncached_decode.int8.onnx"},
				Required: true,
			},
			{
				Role:     "cached_decoder",
				AllOf:    []string{"cached_decode.int8.onnx"},
				Required: true,
			},
			{
				Role:     "tokens",
				AllOf:    []string{"tokens.txt"},
				Required: true,
			},
		},
		ProviderOrder: []string{"cpu"},
		NumThreads:    4,
	},
}

var languageProfiles = map[string]LanguageProfile{
	DefaultLanguageID: {
		ID:             DefaultLanguageID,
		DisplayName:    "Chinese",
		NativeName:     "中文",
		UILocale:       "zh",
		SystemMatchers: []string{"zh", "zh-CN", "zh-Hans", "zh-Hant"},
		DefaultModelID: DefaultModelID,
	},
	EnglishLanguageID: {
		ID:             EnglishLanguageID,
		DisplayName:    "English",
		NativeName:     "English",
		UILocale:       "en",
		SystemMatchers: []string{"en", "en-US", "en-GB", "en-AU", "en-CA"},
		DefaultModelID: MoonshineModelID,
	},
}

func GetModelProfile(id string) (ModelProfile, error) {
	if id == "" {
		id = DefaultModelID
	}
	profile, ok := modelProfiles[id]
	if !ok {
		return ModelProfile{}, fmt.Errorf("unknown model profile: %s", id)
	}
	return profile, nil
}

func DefaultModelProfile() ModelProfile {
	return modelProfiles[DefaultModelID]
}

func ListModelProfiles() []ModelProfile {
	profiles := make([]ModelProfile, 0, len(modelProfiles))
	for _, profile := range modelProfiles {
		profiles = append(profiles, profile)
	}
	return profiles
}

func NormalizeModelID(id string) string {
	if _, err := GetModelProfile(id); err != nil {
		return DefaultModelID
	}
	return id
}

func GetLanguageProfile(id string) (LanguageProfile, error) {
	profile, ok := languageProfiles[id]
	if !ok {
		return LanguageProfile{}, fmt.Errorf("unknown language profile: %s", id)
	}
	return profile, nil
}

func ListLanguageProfiles() []LanguageProfile {
	profiles := make([]LanguageProfile, 0, len(languageProfiles))
	for _, profile := range languageProfiles {
		profiles = append(profiles, profile)
	}
	sort.Slice(profiles, func(i, j int) bool {
		if profiles[i].ID == DefaultLanguageID {
			return true
		}
		if profiles[j].ID == DefaultLanguageID {
			return false
		}
		return profiles[i].ID < profiles[j].ID
	})
	return profiles
}

func NormalizeLanguageID(id string) string {
	if _, err := GetLanguageProfile(id); err != nil {
		return DefaultLanguageID
	}
	return id
}