from __future__ import annotations

import json
from pathlib import Path

from .backends.image import DemoImageBackend
from .backends.text import DemoTextBackend
from .orchestrator import ForestOrchestrator
from .trace import TraceRecorder

SCENARIOS = (
    ("work", "starting a new design job and worried about belonging"),
    ("study", "preparing for a difficult exam after one disappointing result"),
    ("creative", "returning to a novel after months of feeling stuck"),
)


def build_trace_dataset(output: Path) -> int:
    output.unlink(missing_ok=True)
    output.parent.mkdir(parents=True, exist_ok=True)
    row_count = 0
    for index, (category, situation) in enumerate(SCENARIOS):
        trace_id = f"forest-{index + 1:02d}"
        temporary = output.with_name(f".{trace_id}.jsonl")
        temporary.unlink(missing_ok=True)
        orchestrator = ForestOrchestrator(
            text_backend=DemoTextBackend(),
            image_backend=DemoImageBackend(width=128, height=96),
            trace_recorder=TraceRecorder(temporary),
        )
        list(orchestrator.generate("Sample visitor", situation, seed=3407 + index * 10))
        records = [
            {
                "stage": "scenario",
                "scenario_category": category,
                "public_synthetic_input": situation,
            },
            *[
                json.loads(line)
                for line in temporary.read_text(encoding="utf-8").splitlines()
            ],
        ]
        with output.open("a", encoding="utf-8") as trace_file:
            for sequence, record in enumerate(records):
                row = {
                    "trace_id": trace_id,
                    "sequence": sequence,
                    "stage": record["stage"],
                    "model": record.get("model", ""),
                    "payload_json": json.dumps(record, ensure_ascii=True, sort_keys=True),
                }
                trace_file.write(json.dumps(row, ensure_ascii=True) + "\n")
                row_count += 1
        temporary.unlink()
    return row_count