import sys
import os
import time

ACE_ROOT = os.path.join(os.path.dirname(__file__), "ACE-Step-1.5")
sys.path.insert(0, ACE_ROOT)
os.environ["PATH"] = ACE_ROOT + os.pathsep + os.environ.get("PATH", "")

from graph_engine import langgraph_app


def run_production_pipeline(user_session_id: str, raw_story: str):
    config = {"configurable": {"thread_id": user_session_id}}

    for event in langgraph_app.stream({"script": raw_story}, config):
        pass

    final_state = langgraph_app.get_state(config).values
    lyrics = final_state.get("generated_lyrics", "")
    style_tags = final_state.get("music_style_tags", "")

    if not lyrics or not style_tags:
        print("\n[Error]: Failed to generate lyrics or style tags.")
        return

    print(f"\n--- Composed Lyrics ---\n{lyrics}")
    print(f"\n--- Style Tags ---\n{style_tags}")

    print("\n[System]: Initializing ACE-Step 1.5 Diffusion Engine...")

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

    save_dir = os.path.join(os.path.dirname(__file__), "generated_songs")
    os.makedirs(save_dir, exist_ok=True)

    dit_handler = AceStepHandler()
    status_msg, success = dit_handler.initialize_service(
        project_root=ACE_ROOT,
        config_path="acestep-v15-turbo",
        device="auto",
        offload_to_cpu=False,
    )
    if not success:
        print(f"\n[Error]: Failed to initialize ACE-Step: {status_msg}")
        return

    params = GenerationParams(
        task_type="text2music",
        thinking=False,
        caption=style_tags,
        lyrics=lyrics,
        duration=60,
        inference_steps=8,
    )
    gen_config = GenerationConfig(
        batch_size=1,
        audio_format="wav",
    )

    print("\n[System]: Generating audio...")
    result = generate_music(
        dit_handler, None, params=params, config=gen_config, save_dir=save_dir
    )

    if result.success and result.audios:
        audio_path = result.audios[0]["path"]
        print(f"\n[Success]: Track compiled! Saved as {audio_path}")
        print("[System]: Playing track...")
        os.startfile(audio_path)
    else:
        print(f"\n[Error]: Generation failed: {result.status_message}")


if __name__ == "__main__":
    print("=== Local Song Generator ===")
    story_input = input("Enter your story/script for the song: ").strip()
    if not story_input:
        story_input = "A lonely astronaut lost in deep space, watching earth fade away into darkness."
        print(f"[Using default story]: {story_input}")
    session_id = f"session_{int(time.time())}"
    run_production_pipeline(user_session_id=session_id, raw_story=story_input)