from backend import run_llm, generate_variants
from evaluator import final_score


def optimize_prompt(base_prompt: str, reference_answer: str = ""):
    variants = generate_variants(base_prompt)

    batch_prompt = """
You are a helpful AI assistant.

Answer each prompt separately.

Format strictly:
1. <answer>
2. <answer>
3. <answer>
4. <answer>
5. <answer>

Prompts:
"""

    for i, p in enumerate(variants, 1):
        batch_prompt += f"{i}. {p}\n"

    combined_output = run_llm(batch_prompt)

    answers = []
    for i in range(1, len(variants) + 1):
        parts = combined_output.split(f"{i}.")
        if len(parts) > 1:
            answers.append(parts[1].strip())

    results = []
    for p, out in zip(variants, answers):
        score = final_score(out, reference_answer)
        results.append({
            "prompt": p,
            "output": out,
            "score": score
        })

    results.sort(key=lambda x: x["score"], reverse=True)
    return results