"""Format raw qwenapp Aquila access-log trace into the standard release format.

qwenapp traces differ from coding-agent traces in two ways:
  1. Prompts are pre-rendered strings (not a `messages` list).
  2. No tool calls — this is a general chat + RAG workload.

The script reads the raw JSONL, groups rows by session_id, assigns
chat_id / parent_chat_id / turn, tokenizes the prompt to build hash_ids
for prefix-cache replay, and writes the standard format understood by
the replayer and sample_trace.py.

Usage (run on dash1 where the trace and tokenizer live):
    python scripts/format_qwenapp_trace.py \
        --input ~/ali-trace/trace-qwenapp/0529-1000-1200.jsonl \
        --output ~/ali-trace/trace-qwenapp-formatted/0529-1000-1200.jsonl \
        --tokenizer ~/ali-trace/trace_model_meta/Qwen/Qwen3-Coder-480B-A35B-Instruct/tokenizer.json
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from array import array
from collections import defaultdict
from pathlib import Path

from tokenizers import Tokenizer
from tqdm.auto import tqdm


def _block_digest(block: list[int]) -> bytes:
    digest = hashlib.blake2b(digest_size=16)
    digest.update(len(block).to_bytes(4, "little", signed=False))
    digest.update(array("I", block).tobytes())
    return digest.digest()


def _parse_row(raw: dict) -> dict | None:
    """Extract the fields we need from one raw Aquila access-log row."""
    request_params = raw.get("request_params")
    if isinstance(request_params, str):
        request_params = json.loads(request_params)
    response_params = raw.get("response_params")
    if isinstance(response_params, str):
        response_params = json.loads(response_params)

    if not isinstance(request_params, dict) or not isinstance(response_params, dict):
        return None

    inp = request_params.get("payload", {}).get("input", {})
    resp_payload = response_params.get("payload", {})
    usage = resp_payload.get("usage", {})
    resp_header = response_params.get("header", {})

    prompt = inp.get("prompt", "")
    if not prompt:
        return None

    input_tokens = int(usage.get("input_tokens", 0))
    output_tokens = int(usage.get("output_tokens", 0))
    if input_tokens <= 0:
        return None

    status_code = str(resp_header.get("status_code", raw.get("status_code", "")))
    if status_code and status_code != "200" and status_code != "1000":
        return None

    # Timestamp: prefer the `time` field (ms precision), fall back to __time__
    time_str = raw.get("time", "")
    wall_clock_s = _parse_time_to_epoch(time_str) if time_str else None
    if wall_clock_s is None:
        epoch_s = raw.get("__time__")
        wall_clock_s = float(epoch_s) if epoch_s else 0.0

    session_id = str(inp.get("session_id", ""))
    user_id = str((inp.get("source") or {}).get("user_id", ""))
    request_id = str(raw.get("request_id", inp.get("request_id", "")))

    return {
        "session_id": session_id,
        "user_id": user_id,
        "request_id": request_id,
        "wall_clock_s": wall_clock_s,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "prompt": prompt,
        "total_cost_time_ms": int(raw.get("total_cost_time", 0)),
    }


def _parse_time_to_epoch(time_str: str) -> float | None:
    """Parse '2026-05-29 09:59:59.505' (UTC+8) to epoch seconds."""
    import datetime
    try:
        dt = datetime.datetime.strptime(time_str.strip(), "%Y-%m-%d %H:%M:%S.%f")
    except ValueError:
        try:
            dt = datetime.datetime.strptime(time_str.strip(), "%Y-%m-%d %H:%M:%S")
        except ValueError:
            return None
    tz_utc8 = datetime.timezone(datetime.timedelta(hours=8))
    dt = dt.replace(tzinfo=tz_utc8)
    return dt.timestamp()


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--input", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--tokenizer", type=Path, required=True,
                   help="Path to tokenizer.json (Qwen3 family)")
    p.add_argument("--block-size", type=int, default=512)
    p.add_argument("--ttp", action="store_true",
                   help="Also emit time_to_parent_chat field")
    args = p.parse_args()

    print(f"Loading tokenizer from {args.tokenizer} ...")
    tokenizer = Tokenizer.from_file(str(args.tokenizer))

    # --- Pass 1: read all rows, group by session ---
    print(f"Reading {args.input} ...")
    sessions: dict[str, list[dict]] = defaultdict(list)
    skipped = 0
    total = 0

    with args.input.open("r", encoding="utf-8") as fh:
        for line in tqdm(fh, desc="Read", unit="row"):
            total += 1
            raw = json.loads(line)
            parsed = _parse_row(raw)
            if parsed is None:
                skipped += 1
                continue
            sessions[parsed["session_id"]].append(parsed)

    print(f"  Read {total} rows, kept {total - skipped}, skipped {skipped}")
    print(f"  Sessions: {len(sessions)}")

    # --- Sort each session by timestamp and assign turn/chat_id ---
    for rows in sessions.values():
        rows.sort(key=lambda r: r["wall_clock_s"])

    # Assign globally-unique chat_ids
    next_chat_id = 1
    chat_id_map: dict[str, int] = {}  # request_id -> chat_id
    all_rows: list[dict] = []

    for sid, rows in sessions.items():
        prev_chat_id = -1
        prev_end_s = None
        for turn_idx, row in enumerate(rows, start=1):
            cid = next_chat_id
            next_chat_id += 1
            chat_id_map[row["request_id"]] = cid

            ttp = None
            if prev_end_s is not None:
                ttp = row["wall_clock_s"] - prev_end_s
                if ttp < 0:
                    ttp = 0.0

            all_rows.append({
                "session_id": sid,
                "chat_id": cid,
                "parent_chat_id": prev_chat_id,
                "turn": turn_idx,
                "wall_clock_s": row["wall_clock_s"],
                "input_tokens": row["input_tokens"],
                "output_tokens": row["output_tokens"],
                "prompt": row["prompt"],
                "time_to_parent_chat": ttp,
                "total_cost_time_ms": row["total_cost_time_ms"],
            })

            prev_chat_id = cid
            prev_end_s = row["wall_clock_s"] + row["total_cost_time_ms"] / 1000.0

    # Sort globally by timestamp
    all_rows.sort(key=lambda r: r["wall_clock_s"])

    # Re-zero timestamps
    t0 = all_rows[0]["wall_clock_s"] if all_rows else 0.0

    # --- Pass 2: tokenize and hash ---
    print(f"Tokenizing {len(all_rows)} rows (block_size={args.block_size}) ...")
    block_ids_by_digest: dict[str, int] = {}
    next_block_id = 0

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as out:
        for row in tqdm(all_rows, desc="Tokenize+write", unit="row"):
            token_ids = tokenizer.encode(row["prompt"]).ids

            # Build block hashes
            hash_ids = []
            for i in range(0, len(token_ids), args.block_size):
                block = token_ids[i:i + args.block_size]
                digest_hex = _block_digest(block).hex()
                bid = block_ids_by_digest.get(digest_hex)
                if bid is None:
                    bid = next_block_id
                    next_block_id += 1
                    block_ids_by_digest[digest_hex] = bid
                hash_ids.append(bid)

            # Pad hash_ids to match provider input_tokens if needed
            target_blocks = (row["input_tokens"] + args.block_size - 1) // args.block_size
            while len(hash_ids) < target_blocks:
                pad_key = f"pad:{row['chat_id']}:{len(hash_ids)}"
                pad_hex = hashlib.blake2b(pad_key.encode(), digest_size=16).hexdigest()
                bid = block_ids_by_digest.get(pad_hex)
                if bid is None:
                    bid = next_block_id
                    next_block_id += 1
                    block_ids_by_digest[pad_hex] = bid
                hash_ids.append(bid)
            hash_ids = hash_ids[:target_blocks]

            out_row = {
                "chat_id": row["chat_id"],
                "parent_chat_id": row["parent_chat_id"],
                "timestamp": round(row["wall_clock_s"] - t0, 3),
                "input_length": row["input_tokens"],
                "output_length": row["output_tokens"],
                "type": "app",
                "turn": row["turn"],
                "hash_ids": hash_ids,
                "session_id": row["session_id"],
            }
            if args.ttp and row["time_to_parent_chat"] is not None:
                out_row["time_to_parent_chat"] = round(row["time_to_parent_chat"], 3)

            out.write(json.dumps(out_row, ensure_ascii=False) + "\n")

    # --- Summary ---
    n_sessions = len(sessions)
    n_requests = len(all_rows)
    turns_per = [len(v) for v in sessions.values()]
    multi = sum(1 for t in turns_per if t > 1)
    span_s = all_rows[-1]["wall_clock_s"] - t0 if all_rows else 0
    qps = n_requests / span_s if span_s > 0 else 0
    inp_lens = [r["input_tokens"] for r in all_rows]
    out_lens = [r["output_tokens"] for r in all_rows]

    unique_blocks = next_block_id
    shared_blocks = sum(1 for d, bid in block_ids_by_digest.items()
                        if not d.startswith("pad:"))  # rough
    # Count actually shared (appear in >1 request)
    block_ref_count: dict[int, int] = defaultdict(int)
    # Re-scan is expensive; use the hash_ids we just wrote
    # Instead, just report unique count
    print(f"\n=== Summary ===")
    print(f"  Requests: {n_requests}")
    print(f"  Sessions: {n_sessions}")
    print(f"  Multi-turn: {multi} ({multi*100/n_sessions:.1f}%)")
    print(f"  Turns/session: min={min(turns_per)} max={max(turns_per)} avg={sum(turns_per)/len(turns_per):.1f}")
    print(f"  Input tokens: min={min(inp_lens)} max={max(inp_lens)} avg={sum(inp_lens)/len(inp_lens):.0f}")
    print(f"  Output tokens: min={min(out_lens)} max={max(out_lens)} avg={sum(out_lens)/len(out_lens):.0f}")
    print(f"  Trace span: {span_s:.1f}s ({span_s/60:.1f} min)")
    print(f"  QPS: {qps:.2f}")
    print(f"  Unique hash blocks: {unique_blocks}")
    print(f"\nWrote {n_requests} rows to {args.output}")


if __name__ == "__main__":
    main()