#!/usr/bin/env python3
"""
engine_llm.py – LLM Lyric Expander Module (Local Model: microsoft/phi-2)
Generates only: Song Title + 2 verses (8 lines each, rhyming and flowing).
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Tuple

MODEL_NAME = "microsoft/phi-2"

_tokenizer = None
_model = None

def _load_model():
    global _tokenizer, _model
    if _tokenizer is None:
        _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
        _model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.float32,
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )
    return _tokenizer, _model

def run_llm_lyric_expander(brief_text: str) -> Tuple[str, str]:
    """
    Generate only: Song Title + 2 verses (8 lines each, rhyming and flowing).
    No extra text, no sections, no production notes.
    """
    if not brief_text or brief_text.strip() == "":
        return "⚠️ Please provide a creative brief.", "No brief provided."

    try:
        tokenizer, model = _load_model()

        prompt = f"""You are a songwriter. Based on the creative brief below, output ONLY:
1. A song title
2. Two verses of exactly 8 lines each (16 lines total)
Each verse must rhyme and flow. Do not add any extra text, comments, or sections.

Creative Brief:
{brief_text}

Output format:
Title: [Your Title]

[Verse 1]
Line 1
Line 2
Line 3
Line 4
Line 5
Line 6
Line 7
Line 8

[Verse 2]
Line 1
Line 2
Line 3
Line 4
Line 5
Line 6
Line 7
Line 8

Now output exactly that:"""

        inputs = tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = model.generate(
                inputs.input_ids,
                max_new_tokens=512,
                temperature=0.8,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1
            )
        full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # Remove prompt from output
        lyrics = full_output[len(prompt):].strip()
        if not lyrics:
            lyrics = _fallback_lyrics(brief_text)
        return lyrics, "✅ Lyrics generated using microsoft/phi-2 (local model) – no external API needed."

    except Exception as e:
        return _fallback_lyrics(brief_text), f"❌ Local model error: {str(e)}"

def _fallback_lyrics(brief_text: str) -> str:
    return f"""Title: Echoes of the Creative Brief

[Verse 1]
The image pulses, a hidden beat,
Colors and edges, a rhythmic heat.
The prompt whispers, a story untold,
In the machine, the future unfolds.
The seed is planted, the BPM set,
A vision of sound, a path to beget.
The LLM listens, the brief is read,
A symphony born from the image's thread.

[Verse 2]
Let the music rise, let the lyrics flow,
The truth of the image begins to show.
In the resonance, we find the key,
The answer is here, and it's set to be free.
And so the song completes the frame,
A circle closed, a legend named.
The LYGO suite, the final sight,
A journey from sound into light."""