# ref: https://github.com/Aivis-Project/AivisSpeech-Engine/blob/master/run.py

import gc
import uvicorn
from pathlib import Path
import random
import aiohttp
import aiofiles

from voicevox_engine import __version__
from voicevox_engine.aivm_manager import AivmManager
from voicevox_engine.app.application import generate_app
from voicevox_engine.core.core_initializer import MOCK_VER, initialize_cores
from voicevox_engine.engine_manifest import load_manifest
from voicevox_engine.library.library_manager import LibraryManager
from voicevox_engine.logging import LOGGING_CONFIG, logger
from voicevox_engine.preset.preset_manager import PresetManager
from voicevox_engine.setting.model import CorsPolicyMode
from voicevox_engine.setting.setting_manager import USER_SETTING_PATH, SettingHandler
from voicevox_engine.tts_pipeline.song_engine import make_song_engines_from_cores
from voicevox_engine.tts_pipeline.tts_engine import TTSEngineManager
from voicevox_engine.user_dict.user_dict_manager import UserDictionary
from voicevox_engine.utility.path_utility import (
    engine_root,
    get_save_dir,
)
from voicevox_engine.utility.user_agent_utility import generate_user_agent


def start_aivis_speech() -> None:
    """AivisSpeech Engine を実行する"""
    try:
        # multiprocessing.freeze_support()

        # 起動時の可能な限り早い段階で実行結果をキャッシュしておくのが重要
        generate_user_agent("CPU")

        logger.info(f"AivisSpeech Engine version {__version__}")
        logger.info(f"Engine root directory: {engine_root()}")
        logger.info(f"User data directory: {get_save_dir()}")

        # AivmManager を初期化
        aivm_manager = AivmManager(get_save_dir() / "Models")

        # ごく稀に style_bert_vits2_tts_engine.py (が依存する onnxruntime) のインポート自体に失敗し
        # 例外が発生する環境があるようなので、例外をキャッチしてエラーログに出力できるよう、敢えてルーター初期化時にインポートする
        from voicevox_engine.tts_pipeline.style_bert_vits2_tts_engine import (
            StyleBertVITS2TTSEngine,
        )

        # AivisSpeech Engine 独自の StyleBertVITS2TTSEngine を通常の TTSEngine の代わりに利用
        tts_engines = TTSEngineManager()
        tts_engines.register_engine(
            StyleBertVITS2TTSEngine(aivm_manager, use_gpu=False, load_all_models=False),
            MOCK_VER,
        )

        core_manager = initialize_cores(
            use_gpu=False,
            voicelib_dirs=None,
            voicevox_dir=None,
            runtime_dirs=None,
            cpu_num_threads=16,
            enable_mock=True,
            load_all_models=False,
        )
        # tts_engines = make_tts_engines_from_cores(core_manager)
        song_engines = make_song_engines_from_cores(core_manager)
        # assert len(tts_engines.versions()) != 0, "音声合成エンジンがありません。"
        assert len(song_engines.versions()) != 0, "音声合成エンジンがありません。"

        setting_loader = SettingHandler(USER_SETTING_PATH)

        # 複数方式で指定可能な場合、優先度は上から「引数」「環境変数」「設定ファイル」「デフォルト値」

        cors_policy_mode = CorsPolicyMode.all
        allow_origin = ["*"]

        preset_path = get_save_dir() / "presets.yaml"
        preset_manager = PresetManager(preset_path)

        user_dict = UserDictionary()

        engine_manifest = load_manifest(Path("engine/engine_manifest.json"))

        library_manager = LibraryManager(
            # get_save_dir() / "installed_libraries",
            # AivisSpeech では利用しない LibraryManager によるディレクトリ作成を防ぐため、get_save_dir() 直下を指定
            get_save_dir(),
            engine_manifest.supported_vvlib_manifest_version,
            engine_manifest.brand_name,
            engine_manifest.name,
            engine_manifest.uuid,
        )

        root_dir = engine_root()
        character_info_dir = root_dir / "resources" / "character_info"
        # NOTE: ENGINE v0.19 以前向けに後方互換性を確保する
        if not character_info_dir.exists():
            character_info_dir = root_dir / "speaker_info"

        # ASGI に準拠した AivisSpeech Engine アプリケーションを生成する
        app = generate_app(
            tts_engines,
            song_engines,
            aivm_manager,
            core_manager,
            setting_loader,
            preset_manager,
            user_dict,
            engine_manifest,
            library_manager,
            cancellable_engine=None,
            character_info_dir=character_info_dir,
            cors_policy_mode=cors_policy_mode,
            allow_origin=allow_origin,
            disable_mutable_api=False,
        )

        # 起動処理にのみに要したメモリを開放
        gc.collect()

        # AivisSpeech Engine サーバーを起動
        # NOTE: デフォルトは ASGI に準拠した HTTP/1.1 サーバー
        uvicorn.run(app, host="127.0.0.1", port=10101, log_config=LOGGING_CONFIG)

    except Exception as e:
        logger.error("Unexpected error occurred during engine startup:", exc_info=e)
        raise e


def random_str() -> str:
    num = random.randint(10000, 99999)
    return str(num)


async def download_model(model_url: str) -> None:
    save_dir = get_save_dir() / "Models"
    save_dir.mkdir(parents=True, exist_ok=True)

    url = Path(model_url)
    model_id = url.stem
    model_path = save_dir / f"{model_id}.aivmx"

    if model_path.exists():
        logger.info(
            f"Model {model_id} already exists at {model_path}. Skipping download."
        )
        return

    download_url = f"https://api.aivis-project.com/v1/aivm-models/{model_id}/download?model_type=AIVMX"

    logger.info("Downloading model from {download_url} to {model_path}...")

    async with aiohttp.ClientSession() as session:
        try:
            async with session.get(
                download_url,
            ) as res:
                res.raise_for_status()

                # streaming download
                async with aiofiles.open(model_path, "wb") as f:
                    async for chunk in res.content.iter_chunked(1024 * 1024):
                        await f.write(chunk)

                logger.info(f"Model downloaded to {model_path}")

        except Exception as e:
            logger.error(f"Failed to download model: {e}")


if __name__ == "__main__":
    # AivisSpeech Engine を起動
    start_aivis_speech()