import json
import os
import random

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "benchmark", "locomo", "tasks.json")


def generate_conversation(lang: str, num_turns: int = 13):
    personas = {
        "es": ["Juan", "Ana", "Carlos", "María", "Luis", "Sofía"],
        "en": ["Alice", "Bob", "Charlie", "Diana", "Eve", "Frank"],
    }
    topics = {
        "es": [
            "vive en Madrid",
            "es ingeniera de software",
            "trabaja en inteligencia artificial",
            "desarrolla SHCortex",
            "usan Rust",
            "le gusta correr",
            "colecciona sellos",
            "estudia medicina",
            "viajó a Japón",
            "toca piano",
        ],
        "en": [
            "lives in New York",
            "is a data scientist",
            "works on LLMs",
            "builds memory systems",
            "uses Python",
            "enjoys hiking",
            "collects cameras",
            "studies physics",
            "traveled to Japan",
            "plays guitar",
        ],
    }
    names = personas[lang]
    facts = topics[lang]
    random.seed(42)
    conv = []
    used_facts = random.sample(facts, k=min(num_turns // 2, len(facts)))
    for i in range(num_turns):
        speaker = names[i % len(names)]
        fact = used_facts[i % len(used_facts)]
        if lang == "es":
            line = f"{speaker}: Hola, {speaker.lower()} {fact}."
        else:
            line = f"{speaker}: Hi, I'm {speaker} and I {fact}."
        conv.append(line)
    return conv


def generate_questions(lang: str, conversation):
    # Heurística: última entidad nombrada como respuesta
    if lang == "es":
        qs = [
            "¿En qué ciudad vive la primera persona?",
            "¿Cuál es la profesión de la segunda persona?",
            "¿Qué proyecto desarrolla la primera persona?",
            "¿Qué lenguaje de programación usan?",
        ]
    else:
        qs = [
            "Where does the first person live?",
            "What is the second person's profession?",
            "What project is the first person building?",
            "What programming language do they use?",
        ]
    return qs


def main():
    tasks = []
    for i in range(3):
        lang = "es" if i % 2 == 0 else "en"
        conv = generate_conversation(lang, num_turns=13)
        qs = generate_questions(lang, conv)
        tasks.append({
            "id": f"auto_{i+1}",
            "lang": lang,
            "conversation": conv,
            "questions": qs,
            "answers": ["" for _ in qs],
        })
    with open(OUT, "w", encoding="utf-8") as f:
        json.dump(tasks, f, indent=2, ensure_ascii=False)
    print(f"Generated {len(tasks)} tasks -> {OUT}")


if __name__ == "__main__":
    main()