import os
from huggingface_hub import InferenceClient

# ---- TOKEN ----
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise RuntimeError("HF_TOKEN not found. Add it in Space → Settings → Repository secrets")

# ---- MODEL ----
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
client = InferenceClient(model=MODEL_NAME, token=HF_TOKEN)


def run_llm(prompt: str) -> str:
    """Run prompt using conversational API (supported by this model)."""
    try:
        messages = [{"role": "user", "content": prompt}]
        resp = client.chat_completion(
            messages=messages,
            max_tokens=300,
            temperature=0.7,
        )
        return resp.choices[0].message["content"].strip()
    except Exception as e:
        # raise clean error (no stacktrace in UI)
        raise RuntimeError(f"LLM call failed: {e}")


def generate_variants(base_prompt: str):
    meta_prompt = f"""
Rewrite the following prompt in 5 improved ways:
1. More specific
2. More technical
3. More concise
4. More creative
5. Step-by-step style

Return ONLY the rewritten prompts, one per line.

Original prompt:
{base_prompt}
"""
    text = run_llm(meta_prompt)

    variants = []
    for line in text.split("\n"):
        line = line.strip("-• ").strip()
        if line:
            variants.append(line)

    # Fallback if parsing fails
    if len(variants) < 3:
        variants = [
            f"You are an expert. {base_prompt}",
            f"Explain step by step: {base_prompt}",
            f"Give a concise answer: {base_prompt}",
            f"Provide a technical explanation: {base_prompt}",
            f"Explain creatively: {base_prompt}",
        ]

    return variants