#!/usr/bin/env python3
"""Export training data as JSONL for fine-tuning (Well-Tuned badge).

Exports two datasets:
  training_data/entity_generations.jsonl  — entity summoning examples
  training_data/interactions.jsonl        — soul interaction examples
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from world.database import init_database
from world.entities import get_all_entities

OUT_DIR = ROOT / "training_data"

ENTITY_OUT = OUT_DIR / "entity_generations.jsonl"
INTERACTION_OUT = OUT_DIR / "interactions.jsonl"


def export_entities() -> int:
    entities = get_all_entities(limit=1000)
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    count = 0
    with ENTITY_OUT.open("w", encoding="utf-8") as f:
        for e in entities:
            if not e.get("input_description"):
                continue
            from world.locations import get_location_by_id
            loc = get_location_by_id(e["location_id"])
            label = {
                "name": e["name"],
                "display_name": e["display_name"],
                "type": e["type"],
                "appearance": e["appearance"],
                "backstory": e["backstory"],
                "personality_traits": e["personality_traits"],
                "primary_goal": e["primary_goal"],
                "secondary_goal": e["secondary_goal"],
                "primary_fear": e["primary_fear"],
                "speech_style": e["speech_style"],
                "greeting": e["greeting"],
                "suggested_location": loc["name"] if loc else "",
                "arrival_note": e.get("arrival_note", ""),
            }
            f.write(json.dumps({
                "messages": [
                    {"role": "system", "content": "You are the Oracle of Aether Garden. Generate a soul from the visitor's description."},
                    {"role": "user",   "content": f'Summon: "{e["input_description"]}"'},
                    {"role": "assistant", "content": json.dumps(label, ensure_ascii=False)},
                ]
            }, ensure_ascii=False) + "\n")
            count += 1
    return count


def export_interactions() -> int:
    from world.database import db_session
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    count = 0
    with db_session() as conn:
        rows = conn.execute("""
            SELECT i.*,
                   ea.display_name as a_name, ea.type as a_type,
                   ea.appearance as a_appearance, ea.personality_traits as a_traits,
                   ea.primary_goal as a_goal, ea.primary_fear as a_fear,
                   ea.speech_style as a_speech, ea.memory_summary as a_memory,
                   eb.display_name as b_name, eb.type as b_type,
                   eb.appearance as b_appearance, eb.personality_traits as b_traits,
                   eb.primary_goal as b_goal, eb.primary_fear as b_fear,
                   eb.speech_style as b_speech, eb.memory_summary as b_memory,
                   l.name as location_name
            FROM interactions i
            LEFT JOIN entities ea ON ea.id = i.entity_a_id
            LEFT JOIN entities eb ON eb.id = i.entity_b_id
            LEFT JOIN locations l ON l.id = i.location_id
            WHERE i.description IS NOT NULL AND length(i.description) > 30
            LIMIT 2000
        """).fetchall()

    with INTERACTION_OUT.open("w", encoding="utf-8") as f:
        for r in rows:
            r = dict(r)
            if not r.get("a_name") or not r.get("b_name"):
                continue
            context = (
                f"LOCATION: {r.get('location_name','Unknown')}\n"
                f"ENTITY A: {r['a_name']} ({r['a_type']})\n"
                f"  Goal: {r.get('a_goal','')}\n  Fear: {r.get('a_fear','')}\n"
                f"  Memory: {r.get('a_memory','')}\n"
                f"ENTITY B: {r['b_name']} ({r['b_type']})\n"
                f"  Goal: {r.get('b_goal','')}\n  Fear: {r.get('b_fear','')}\n"
                f"  Memory: {r.get('b_memory','')}"
            )
            output = {
                "interaction_type": r.get("interaction_type","chance_meeting"),
                "description": r.get("description",""),
                "book_of_ages_entry": r.get("book_of_ages_entry",""),
            }
            f.write(json.dumps({
                "messages": [
                    {"role": "system", "content": "You are the simulation engine of Aether Garden. Generate what happened between two souls."},
                    {"role": "user",   "content": context},
                    {"role": "assistant", "content": json.dumps(output, ensure_ascii=False)},
                ]
            }, ensure_ascii=False) + "\n")
            count += 1
    return count


def main():
    init_database()
    n_entities = export_entities()
    n_ints = export_interactions()
    total = n_entities + n_ints
    print(f"Exported {n_entities} entity examples  →  {ENTITY_OUT}")
    print(f"Exported {n_ints} interaction examples →  {INTERACTION_OUT}")
    print(f"Total: {total} training examples")
    print()
    print("To fine-tune on HuggingFace AutoTrain:")
    print("  1. Upload both JSONL files to a private HF dataset")
    print("  2. Run AutoTrain with 'chat' task type on Qwen2.5-3B-Instruct")
    print("  3. Set FINE_TUNED_MODEL env var in your Modal deployment")
    print("  4. Redeploy: modal deploy modal_app.py")


if __name__ == "__main__":
    main()