from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class TraceRequest:
    request_id: str
    session_id: str
    chat_id: int
    parent_chat_id: int
    timestamp_s: float
    input_length: int
    output_length: int
    request_type: str
    turn_id: int
    hash_ids: tuple[int, ...]


def load_trace(path: Path, *, request_limit: int | None = None) -> list[TraceRequest]:
    chat_to_session: dict[int, str] = {}
    requests: list[TraceRequest] = []

    with path.open("r", encoding="utf-8") as handle:
        for index, line in enumerate(handle):
            if request_limit is not None and len(requests) >= request_limit:
                break

            payload = json.loads(line)
            chat_id = int(payload["chat_id"])
            parent_chat_id = int(payload["parent_chat_id"])
            session_id = _resolve_session_id(
                chat_id=chat_id,
                parent_chat_id=parent_chat_id,
                chat_to_session=chat_to_session,
            )
            turn_id = int(payload["turn"])
            request_id = f"{session_id}:{turn_id}:{chat_id}:{index}"
            requests.append(
                TraceRequest(
                    request_id=request_id,
                    session_id=session_id,
                    chat_id=chat_id,
                    parent_chat_id=parent_chat_id,
                    timestamp_s=float(payload["timestamp"]),
                    input_length=int(payload["input_length"]),
                    output_length=int(payload["output_length"]),
                    request_type=str(payload["type"]),
                    turn_id=turn_id,
                    hash_ids=tuple(int(item) for item in payload.get("hash_ids", [])),
                )
            )

    return requests


def build_synthetic_prompt(
    request: TraceRequest,
    *,
    block_token_budget: int = 24,
) -> str:
    return " ".join(build_synthetic_prompt_tokens(request, block_token_budget=block_token_budget))


def build_synthetic_prompt_tokens(
    request: TraceRequest,
    *,
    block_token_budget: int = 24,
) -> list[str]:
    tokens: list[str] = []
    for hash_id in request.hash_ids:
        for offset in range(block_token_budget):
            tokens.append(f"blk{hash_id}_{offset}")

    while len(tokens) < request.input_length:
        tokens.append(f"fill_{len(tokens) % 64}")

    return tokens[: request.input_length]


def build_synthetic_append_chunk(
    request: TraceRequest,
    append_length: int,
) -> str:
    if append_length <= 0:
        return ""
    return " ".join(
        f"turn{request.turn_id}_append_{request.chat_id}_{offset}"
        for offset in range(append_length)
    )


def _resolve_session_id(
    *,
    chat_id: int,
    parent_chat_id: int,
    chat_to_session: dict[int, str],
) -> str:
    if parent_chat_id < 0:
        session_id = str(chat_id)
    else:
        session_id = chat_to_session.get(parent_chat_id, str(parent_chat_id))
    chat_to_session[chat_id] = session_id
    return session_id
