from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from math import ceil
from pathlib import Path


BLOCK_TOKEN_BUDGET = 24


@dataclass(frozen=True)
class SmallAppendTraceConfig:
    output_path: Path
    session_count: int = 8
    turns_per_session: int = 3
    initial_input_length: int = 10_000
    append_input_length: int = 1_000
    output_length: int = 1_000
    inter_turn_gap_s: float = 1.0
    session_stagger_s: float = 0.1
    request_type: str = "coder"


@dataclass(frozen=True)
class SmallAppendTraceSummary:
    output_path: str
    session_count: int
    turns_per_session: int
    request_count: int
    initial_input_length: int
    append_input_length: int
    output_length: int
    inter_turn_gap_s: float
    session_stagger_s: float


def write_small_append_trace(config: SmallAppendTraceConfig) -> SmallAppendTraceSummary:
    if config.session_count <= 0:
        raise ValueError("session_count must be > 0")
    if config.turns_per_session <= 0:
        raise ValueError("turns_per_session must be > 0")
    if config.initial_input_length < 0:
        raise ValueError("initial_input_length must be >= 0")
    if config.append_input_length < 0:
        raise ValueError("append_input_length must be >= 0")
    if config.output_length < 0:
        raise ValueError("output_length must be >= 0")

    config.output_path.parent.mkdir(parents=True, exist_ok=True)
    records: list[dict[str, object]] = []
    next_chat_id = 1_000_000

    for session_idx in range(config.session_count):
        root_chat_id = next_chat_id
        previous_chat_id = -1
        session_base_time = session_idx * config.session_stagger_s
        base_block_count = ceil(config.initial_input_length / BLOCK_TOKEN_BUDGET)
        base_hash_ids = [
            _hash_id_for(session_idx=session_idx, block_idx=block_idx)
            for block_idx in range(base_block_count)
        ]

        for turn_idx in range(config.turns_per_session):
            chat_id = root_chat_id if turn_idx == 0 else next_chat_id
            if turn_idx > 0:
                next_chat_id += 1

            input_length = config.initial_input_length + turn_idx * (
                config.append_input_length + config.output_length
            )
            total_block_count = ceil(input_length / BLOCK_TOKEN_BUDGET)
            hash_ids = base_hash_ids + [
                _hash_id_for(
                    session_idx=session_idx,
                    block_idx=base_block_count + append_block_idx,
                )
                for append_block_idx in range(max(0, total_block_count - base_block_count))
            ]

            records.append(
                {
                    "chat_id": chat_id,
                    "parent_chat_id": previous_chat_id,
                    "timestamp": session_base_time
                    + turn_idx * config.inter_turn_gap_s,
                    "input_length": input_length,
                    "output_length": config.output_length,
                    "type": config.request_type,
                    "turn": turn_idx + 1,
                    "hash_ids": hash_ids,
                }
            )
            previous_chat_id = chat_id

        next_chat_id += 1

    records.sort(key=lambda item: float(item["timestamp"]))
    with config.output_path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, sort_keys=True) + "\n")

    summary = SmallAppendTraceSummary(
        output_path=str(config.output_path),
        session_count=config.session_count,
        turns_per_session=config.turns_per_session,
        request_count=len(records),
        initial_input_length=config.initial_input_length,
        append_input_length=config.append_input_length,
        output_length=config.output_length,
        inter_turn_gap_s=config.inter_turn_gap_s,
        session_stagger_s=config.session_stagger_s,
    )
    summary_path = config.output_path.with_suffix(
        config.output_path.suffix + ".summary.json"
    )
    with summary_path.open("w", encoding="utf-8") as handle:
        json.dump(asdict(summary), handle, indent=2, sort_keys=True)
    return summary


def _hash_id_for(*, session_idx: int, block_idx: int) -> int:
    return session_idx * 1_000_000 + block_idx