#!/usr/bin/env python3
"""XL turbo (4B DiT) — новый вдохновляющий трек для Ника."""
import os
import sys
import time

for v in ['http_proxy', 'https_proxy', 'HTTP_PROXY', 'HTTPS_PROXY', 'ALL_PROXY']:
    os.environ.pop(v, None)

PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, PROJECT_ROOT)

from loguru import logger
from acestep.handler import AceStepHandler
from acestep.inference import GenerationParams, GenerationConfig, generate_music

SAVE_DIR = os.path.join(PROJECT_ROOT, "output", "personal_new")

CAPTION = (
    "An uplifting, cinematic pop-rock anthem with a powerful male vocal. "
    "Building drums, bright synth pads, soaring electric guitar, emotional "
    "and hopeful mood. Anthem-like chorus, sense of a new dawn and forward "
    "motion, determined and inspiring."
)

LYRICS = """[Intro]

[Verse]
Рассвет на стекле, в окне уже не тьма
Я снова здесь, я снова в борьбе
Прошлые ошибки — лишь следы на пути
Я не сдамся, я иду вперёд

[Pre-Chorus]
Пусть руки дрожат, но в груди огонь
Я не сломан — я только в пути

[Chorus]
Восход, восход — я иду навстречу свету
Восход, восход — хоть и трудно порой
Я ещё здесь, я ещё в строю
И мой путь только начинается

[Verse]
Дорога длинная, впереди много гор
Но каждый восход делает сильнее меня
Я не один, я веду этот бой
За себя и за тех, кто со мной

[Chorus]
Восход, восход — я иду навстречу свету
Восход, восход — хоть и трудно порой
Я ещё здесь, я ещё в строю
И мой путь только начинается

[Bridge]
Не бойся рассвета, он снова придет
Не бойся пути — он твой

[Chorus]
Восход, восход — я иду навстречу свету
Восход, восход — хоть и трудно порой
Я ещё здесь, я ещё в строю
И мой путь только начинается

[Outro]"""


def main():
    os.makedirs(SAVE_DIR, exist_ok=True)

    logger.info("Инициализация DiT handler (acestep-v15-xl-turbo, cuda:0)...")
    t0 = time.time()
    dit = AceStepHandler()
    status, ok = dit.initialize_service(
        project_root=PROJECT_ROOT,
        config_path="acestep-v15-xl-turbo",
        device="cuda",
        offload_to_cpu=True,
    )
    if not ok:
        logger.error(f"DiT init failed: {status}")
        sys.exit(1)
    logger.info(f"DiT loaded за {time.time() - t0:.1f}s — {status}")

    params = GenerationParams(
        task_type="text2music",
        thinking=False,
        caption=CAPTION,
        lyrics=LYRICS,
        vocal_language="ru",
        bpm=128,
        keyscale="F major",
        timesignature="",
        duration=150,
        inference_steps=8,
        guidance_scale=1.0,
        seed=-1,
    )
    config = GenerationConfig(batch_size=1, audio_format="wav")

    logger.info("Генерация на GPU (замер времени)...")
    t0 = time.time()
    result = generate_music(dit, None, params, config, SAVE_DIR)
    elapsed = time.time() - t0

    if result.success:
        logger.info(f"OK за {elapsed:.1f}s")
        for a in result.audios:
            logger.info(f"  -> {a.get('path')}")
    else:
        logger.error(f"FAILED за {elapsed:.1f}s — {result.status_message}")
        sys.exit(1)


if __name__ == "__main__":
    main()