from transformers import AutoModelForCausalLM, AutoTokenizer


_model = None
_tokenizer = None


def _load_qwen():
    global _model, _tokenizer
    if _model is None:
        model_id = "Qwen/Qwen3-8B"
        _tokenizer = AutoTokenizer.from_pretrained(model_id)
        _model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype="auto",
            device_map="auto"
        )
    return _model, _tokenizer


def continue_lyrics(existing_lyrics, key=None, bpm=None, style_hint=None, num_lines=8):
    """
    takes existing lyrics and writes more in the same style.
    key/bpm/style_hint give the model musical context.
    """
    model, tokenizer = _load_qwen()

    context_parts = []
    if key:
        context_parts.append(f"The song is in {key}")
    if bpm:
        context_parts.append(f"at {bpm} BPM")
    if style_hint:
        context_parts.append(f"with a {style_hint} feel")

    context = ", ".join(context_parts) + "." if context_parts else ""

    prompt = f"""You are a songwriter. Continue the following lyrics naturally,
matching the tone, rhythm, and imagery. Write exactly {num_lines} new lines.
Do not repeat existing lines. Do not add commentary or explanations.
{context}

Existing lyrics:
{existing_lyrics}

Continuation:"""

    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    inputs = tokenizer([text], return_tensors="pt").to(model.device)

    output = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.8,
        top_p=0.9,
        do_sample=True
    )

    generated = output[0][inputs.input_ids.shape[1]:]
    result = tokenizer.decode(generated, skip_special_tokens=True)

    # trim to requested line count
    lines = [l for l in result.strip().split('\n') if l.strip()]
    return '\n'.join(lines[:num_lines])