#!/usr/bin/env python3
"""Animily Music - 메인 스케줄러

1시간 / 2시간 / 12시간(1h+2h 혼합×4) 영상 생성 → 유튜브 업로드.
GPU 메모리 관리: ACE-Step → 종료 → FLUX → 종료 → ffmpeg.
"""
import gc
import os
import signal
import subprocess
import sys
import time
from datetime import datetime

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

from config import OUTPUT_DIR, LOG_DIR, SEGMENTS_FOR_1H, SEGMENTS_FOR_2H
from generate_music import generate_2h_music, _wait_for_acestep
from generate_image import generate_image, load_image_prompt, _stop_comfyui, generate_thumbnail
from render_video import render_video, concat_videos
from upload_youtube import upload_video, create_or_get_playlist, add_to_playlist


def _kill_acestep():
    """ACE-Step 서버 관련 프로세스 GPU 메모리 해제 확인"""
    # ACE-Step은 외부 서버이므로 직접 종료하지 않음
    # 다만 GPU 메모리가 해제될 때까지 대기
    print("  [GPU] ACE-Step 작업 완료, GPU 메모리 해제 대기...", flush=True)
    time.sleep(10)
    gc.collect()


def _cleanup_outputs(*paths):
    """임시 파일 정리"""
    for p in paths:
        if p and os.path.exists(p):
            os.remove(p)
            print(f"  [CLEAN] 삭제: {os.path.basename(p)}", flush=True)


def generate_1h_music(animal_type="dog", style_index=0):
    """1시간 음악 생성"""
    from generate_music import (
        _load_prompts, _submit_task, _poll_result,
        _download_audio, _crossfade_segments,
    )
    from config import SEGMENT_DURATION, CROSSFADE_SEC

    prompts = _load_prompts(animal_type)
    style = prompts["styles"][style_index % len(prompts["styles"])]

    base_caption = style["caption"]
    bpm = style["bpm"]
    keyscale = style["keyscale"]
    variations = style.get("variations", [])

    segments_needed = SEGMENTS_FOR_1H  # 12 * 295s ≈ 59분

    print(f"\n{'='*60}")
    print(f"[MUSIC] {animal_type} 1시간 음악 생성 ({segments_needed}세그먼트)")
    print(f"{'='*60}\n")

    segment_paths = []
    for i in range(segments_needed):
        seg_path = os.path.join(OUTPUT_DIR, f"seg_1h_{animal_type}_{i:02d}.wav")

        if variations:
            var = variations[i % len(variations)]
            caption = f"{base_caption} {var}"
        else:
            caption = base_caption

        print(f"  [SEG {i+1}/{segments_needed}] 생성 중...", flush=True)
        t0 = time.time()
        try:
            task_id = _submit_task(caption, bpm, keyscale, SEGMENT_DURATION)
            audio_info = _poll_result(task_id, timeout=900)
            _download_audio(audio_info, seg_path)
            print(f"  [SEG {i+1}/{segments_needed}] 완료 ({time.time()-t0:.0f}초)")
            segment_paths.append(seg_path)
        except Exception as e:
            print(f"  [SEG {i+1}] 실패: {e}")
            continue

    if len(segment_paths) < 6:
        raise RuntimeError(f"1시간 음악: {len(segment_paths)}개만 성공 — 최소 6개 필요")

    output_path = os.path.join(OUTPUT_DIR, f"music_{animal_type}_1h.wav")
    _crossfade_segments(segment_paths, output_path, CROSSFADE_SEC * 1000)

    for p in segment_paths:
        if os.path.exists(p):
            os.remove(p)

    return output_path


def run_pipeline():
    """전체 파이프라인 실행

    생성할 영상:
    1. 1시간짜리 (강아지 솔로 피아노)
    2. 2시간짜리 (강아지 소프트 레게)
    3. 12시간짜리 = (1시간 + 2시간) × 4 반복
    """
    start_time = time.time()
    today = datetime.now().strftime("%Y%m%d")

    print(f"\n{'#'*60}")
    print(f"# Animily Music Pipeline - {today}")
    print(f"{'#'*60}\n")

    results = []
    video_1h_path = None
    video_2h_path = None

    # ============================================================
    # STEP 1: ACE-Step 음악 생성 (1시간 + 2시간)
    # ============================================================
    print("\n[STEP 1] 음악 생성 (ACE-Step)")
    print("=" * 40)

    if not _wait_for_acestep(timeout=30):
        # ACE-Step 서버 시작 시도
        print("  [MUSIC] ACE-Step 서버 시작 중...", flush=True)
        subprocess.Popen(
            ["bash", "-c",
             "cd /home/javamon/ACE-Step-1.5 && source venv/bin/activate && "
             "python -m uvicorn acestep.api_server:app --host 0.0.0.0 --port 8001 --workers 1"],
            stdout=open(os.path.join(LOG_DIR, "acestep.log"), "a"),
            stderr=open(os.path.join(LOG_DIR, "acestep.log"), "a"),
        )
        if not _wait_for_acestep(timeout=180):
            print("  [ERROR] ACE-Step 서버 시작 실패!", flush=True)
            return

    # 1시간 음악 (강아지 - 솔로 피아노, style_index=0)
    music_1h = generate_1h_music("dog", style_index=0)
    print(f"  [MUSIC] 1시간 음악 완료: {music_1h}")

    # 2시간 음악 (강아지 - 소프트 레게, style_index=1)
    music_2h = generate_2h_music("dog", style_index=1)
    print(f"  [MUSIC] 2시간 음악 완료: {music_2h}")

    # ACE-Step 작업 완료 → GPU 해제 대기
    _kill_acestep()

    # ============================================================
    # STEP 2: FLUX 이미지 생성 (영상당 1장)
    # ============================================================
    print("\n[STEP 2] 이미지 생성 (FLUX)")
    print("=" * 40)

    image_1h = os.path.join(OUTPUT_DIR, f"image_1h_{today}.png")
    image_2h = os.path.join(OUTPUT_DIR, f"image_2h_{today}.png")
    thumb_1h = os.path.join(OUTPUT_DIR, f"thumb_1h_{today}.jpg")
    thumb_2h = os.path.join(OUTPUT_DIR, f"thumb_2h_{today}.jpg")

    # 1시간용 이미지
    prompt_1h = load_image_prompt("dog")
    if not generate_image(prompt_1h, image_1h):
        print("  [WARN] 1시간 이미지 생성 실패, 2시간 이미지 공유", flush=True)

    # 2시간용 이미지 (다른 프롬프트)
    prompt_2h = load_image_prompt("dog")
    if not generate_image(prompt_2h, image_2h):
        # 1시간 이미지를 재사용
        if os.path.exists(image_1h):
            subprocess.run(["cp", image_1h, image_2h])

    # 썸네일 생성
    if os.path.exists(image_1h):
        generate_thumbnail(image_1h, thumb_1h)
    if os.path.exists(image_2h):
        generate_thumbnail(image_2h, thumb_2h)

    # FLUX/ComfyUI 종료 → GPU 메모리 해제
    _stop_comfyui()

    # ============================================================
    # STEP 3: 영상 렌더링 (ffmpeg)
    # ============================================================
    print("\n[STEP 3] 영상 렌더링")
    print("=" * 40)

    # 사용할 이미지 결정
    img_for_1h = image_1h if os.path.exists(image_1h) else image_2h
    img_for_2h = image_2h if os.path.exists(image_2h) else image_1h

    # 1시간 영상
    video_1h_path = os.path.join(OUTPUT_DIR, f"pet_music_1h_{today}.mp4")
    if render_video(img_for_1h, music_1h, video_1h_path):
        print(f"  [VIDEO] 1시간 영상 완료")
    else:
        print(f"  [ERROR] 1시간 영상 렌더 실패!")

    # 2시간 영상
    video_2h_path = os.path.join(OUTPUT_DIR, f"pet_music_2h_{today}.mp4")
    if render_video(img_for_2h, music_2h, video_2h_path):
        print(f"  [VIDEO] 2시간 영상 완료")
    else:
        print(f"  [ERROR] 2시간 영상 렌더 실패!")

    # 12시간 영상 = (1h + 2h) × 4 반복
    video_12h_path = os.path.join(OUTPUT_DIR, f"pet_music_12h_{today}.mp4")
    if video_1h_path and video_2h_path and os.path.exists(video_1h_path) and os.path.exists(video_2h_path):
        # 1h + 2h = 3h, × 4 = 12h
        mix_list = [video_1h_path, video_2h_path] * 4
        if concat_videos(mix_list, video_12h_path):
            print(f"  [VIDEO] 12시간 영상 완료")
        else:
            print(f"  [ERROR] 12시간 영상 렌더 실패!")

    # ============================================================
    # STEP 4: YouTube 업로드
    # ============================================================
    print("\n[STEP 4] YouTube 업로드")
    print("=" * 40)

    playlist_id = create_or_get_playlist()

    # 1시간 업로드
    if video_1h_path and os.path.exists(video_1h_path):
        title_1h = "🐕 강아지가 좋아하는 음악 1시간 | 솔로 피아노 수면음악 [과학적 검증]"
        vid_id = upload_video(
            video_1h_path, title_1h,
            thumbnail_path=thumb_1h if os.path.exists(thumb_1h) else None,
            extra_tags=["1시간", "솔로피아노", "수면"],
        )
        if vid_id and playlist_id:
            add_to_playlist(vid_id, playlist_id)
        results.append(("1시간", vid_id))

    # 2시간 업로드
    if video_2h_path and os.path.exists(video_2h_path):
        title_2h = "🐕 강아지가 좋아하는 음악 2시간 | 소프트 레게, 분리불안 완화 [과학적 검증]"
        vid_id = upload_video(
            video_2h_path, title_2h,
            thumbnail_path=thumb_2h if os.path.exists(thumb_2h) else None,
            extra_tags=["2시간", "레게", "분리불안"],
        )
        if vid_id and playlist_id:
            add_to_playlist(vid_id, playlist_id)
        results.append(("2시간", vid_id))

    # 12시간 업로드
    if os.path.exists(video_12h_path):
        title_12h = "🐕 강아지가 좋아하는 음악 12시간 | 수면, 분리불안, 스트레스 해소 [과학적 검증]"
        vid_id = upload_video(
            video_12h_path, title_12h,
            thumbnail_path=thumb_2h if os.path.exists(thumb_2h) else None,
            extra_tags=["12시간", "수면음악", "분리불안", "장시간"],
        )
        if vid_id and playlist_id:
            add_to_playlist(vid_id, playlist_id)
        results.append(("12시간", vid_id))

    # ============================================================
    # STEP 5: 정리
    # ============================================================
    print("\n[STEP 5] 정리")
    print("=" * 40)

    # 임시 음악 파일 삭제 (영상에 포함됐으므로)
    _cleanup_outputs(music_1h, music_2h, image_1h, image_2h, thumb_1h, thumb_2h)
    # 영상 파일은 업로드 확인 후 삭제
    for label, vid_id in results:
        if vid_id:
            print(f"  [{label}] 업로드 성공 → 영상 파일 삭제")
    # 업로드 성공한 영상만 삭제
    if results:
        for path in [video_1h_path, video_2h_path, video_12h_path]:
            if path and os.path.exists(path):
                os.remove(path)
                print(f"  [CLEAN] {os.path.basename(path)} 삭제")

    # GPU 메모리 최종 해제 확인
    gc.collect()

    elapsed = time.time() - start_time
    print(f"\n{'#'*60}")
    print(f"# 파이프라인 완료 ({elapsed/60:.1f}분)")
    print(f"# 결과:")
    for label, vid_id in results:
        url = f"https://youtu.be/{vid_id}" if vid_id else "실패"
        print(f"#   {label}: {url}")
    print(f"{'#'*60}\n")


if __name__ == "__main__":
    run_pipeline()