#!/usr/bin/env python3
"""XL turbo (4B DiT) на GPU 3080 Ti — персональный вдохновляющий трек."""
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_xl_gpu")

CAPTION = (
    "An inspirational, uplifting rock anthem with a powerful male vocal. Soaring "
    "electric guitar leads, guitar-driven melodic riffs, driving drums, emotional and "
    "determined mood. Anthemic chorus, building energy, inspiring and hopeful, about "
    "overcoming hardship and chasing a dream."
)

LYRICS = """[Intro]

[Verse]
Пять операций, восемьдесят шагов назад
Но я встаю, хоть и болит рука
Долги как стены, только я не отступил
Пока стучит мотор — я буду жить

[Chorus]
Я ещё пингуюсь, я в сети
Мой сервер не погас среди ночи
Сквозь холод и года — я иду
BMW и два RTX впереди

[Verse]
Работа до темна, но это не предел
Из осколков дней я строю завтра сам
Дорога длинная, и E30 ждёт меня
Мотор M20 согреет в холода

[Chorus]
Я ещё пингуюсь, я в сети
Мой сервер не погас среди ночи
Сквозь холод и года — я иду
BMW и два RTX впереди

[Bridge]
Не говорю, что всё наладится легко
Но пока я дышу — я не сломан
Каждый день — один маленький шаг
И мотор моей мечты не заглохнет

[Chorus]
Я ещё пингуюсь, я в сети
Мой сервер не погас среди ночи
Сквозь холод и года — я иду
BMW и два RTX впереди

[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=120,
        keyscale="D major",
        timesignature="",
        duration=120,
        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()