from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
from typing import Any
from dataclasses import dataclass


REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_LEGACY_SOURCE = Path(
    "/home/admin/cpfs/wjh/bailian-trace/qwen-trace-260311-260317-formatted"
)
DEFAULT_THINKING_SOURCE = Path(
    "/home/admin/cpfs/wjh/bailian-trace/qwen-trace-260321-260327-formatted"
)
DEFAULT_OUTPUT_ROOT = REPO_ROOT / "trace_windows"
LEGACY_TARGET_DATES = [
    "2026-03-11",
    "2026-03-12",
    "2026-03-13",
    "2026-03-14",
    "2026-03-15",
    "2026-03-16",
    "2026-03-17",
]
THINKING_WINDOWS = [
    ("2026-03-21", "1000"),
    ("2026-03-22", "1000"),
    ("2026-03-23", "1000"),
    ("2026-03-24", "1000"),
    ("2026-03-25", "1000"),
    ("2026-03-26", "1000"),
    ("2026-03-27", "1000"),
]
WINDOW_SPECS = {
    "1000": {
        "start_hour": 9,
        "window_start": 3600.0,
        "window_end": 4200.0,
        "slot_label": "10:00-10:10",
        "slot_token": "1000",
    },
    "2200": {
        "start_hour": 21,
        "window_start": 3600.0,
        "window_end": 4200.0,
        "slot_label": "22:00-22:10",
        "slot_token": "2200",
    },
}
TRACE_SPECS = {
    "chat": {
        "block_size": 64,
        "source_key": "legacy",
        "target_dates": LEGACY_TARGET_DATES,
        "slot_tokens": ("1000", "2200"),
    },
    "coder": {
        "block_size": 512,
        "source_key": "legacy",
        "target_dates": LEGACY_TARGET_DATES,
        "slot_tokens": ("1000", "2200"),
    },
    "thinking": {
        "block_size": 64,
        "source_key": "thinking",
        "date_slot_tokens": THINKING_WINDOWS,
    },
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Prepare canonical trace windows under the current aituner repo by reading the "
            "raw formatted Bailian trace directories and merging prompt payloads."
        )
    )
    parser.add_argument("--legacy-source", type=Path, default=DEFAULT_LEGACY_SOURCE)
    parser.add_argument("--thinking-source", type=Path, default=DEFAULT_THINKING_SOURCE)
    parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT)
    parser.add_argument(
        "--workloads",
        default="chat,coder,thinking",
        help="Comma-separated workloads to materialize. Defaults to chat,coder,thinking.",
    )
    parser.add_argument("--sample-seed", type=int, default=20260325)
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def resolve_session_root(row: dict[str, Any], root_of: dict[Any, Any]) -> Any:
    """Resolve the session root chat_id for a trace row.

    Sessions are multi-turn chains linked via parent_chat_id (turn>1 points to the
    parent turn's chat_id, the root turn has parent_chat_id=-1). Because parent
    turns precede their children in time, a single streaming pass that records
    chat_id -> root resolves the full chain. Rows whose parent is not yet known
    (e.g. it fell outside the materialized span) fall back to the parent id so
    siblings still group together.
    """
    chat_id = row.get("chat_id")
    parent = row.get("parent_chat_id")
    parent_is_root = (
        parent is None
        or (isinstance(parent, (int, float)) and not isinstance(parent, bool) and int(parent) < 0)
    )
    root = chat_id if parent_is_root else root_of.get(parent, parent)
    if chat_id is not None:
        root_of[chat_id] = root
    return root


def session_uniform(*, seed: int, window_id: str, session_root: Any) -> float:
    """Deterministic per-session uniform score in [0, 1).

    All turns of a session share one score, so thresholding sampling_u keeps or
    drops whole sessions and preserves intra-session prefix (KV-cache) reuse.
    """
    payload = json.dumps(
        {
            "seed": seed,
            "window_id": window_id,
            "session_root": session_root,
        },
        sort_keys=True,
        separators=(",", ":"),
    ).encode("utf-8")
    digest = hashlib.blake2b(payload, digest_size=8).digest()
    return int.from_bytes(digest, byteorder="big", signed=False) / float(1 << 64)


def build_source_filename(*, trace_type: str, date_text: str, slot_token: str) -> str:
    month, day = date_text[5:7], date_text[8:10]
    start_hour = int(WINDOW_SPECS[slot_token]["start_hour"])
    return (
        f"qwen_{trace_type}_blksz_{TRACE_SPECS[trace_type]['block_size']}_"
        f"{month}{day}{start_hour:02d}-{month}{day}{start_hour + 2:02d}.jsonl"
    )


def parse_workloads(text: str) -> tuple[str, ...]:
    workloads = tuple(part.strip() for part in str(text).split(",") if part.strip())
    if not workloads:
        raise ValueError("--workloads must contain at least one workload")
    unknown = [item for item in workloads if item not in TRACE_SPECS]
    if unknown:
        raise ValueError(f"Unknown workloads: {unknown}")
    return workloads


def build_windows(
    source_dirs: dict[str, Path], *, workloads: tuple[str, ...]
) -> list[dict[str, Any]]:
    windows: list[dict[str, Any]] = []
    for trace_type in workloads:
        spec = TRACE_SPECS[trace_type]
        source_dir = source_dirs[str(spec["source_key"])]
        if "date_slot_tokens" in spec:
            date_slot_tokens = list(spec["date_slot_tokens"])
        else:
            date_slot_tokens = [
                (date_text, slot_token)
                for slot_token in spec["slot_tokens"]
                for date_text in spec["target_dates"]
            ]
        sample_order_by_group: dict[str, int] = {}
        for date_text, slot_token in date_slot_tokens:
            window_spec = WINDOW_SPECS[slot_token]
            date_token = date_text.replace("-", "")
            sample_order = sample_order_by_group.setdefault(slot_token, 0)
            sample_order_by_group[slot_token] += 1
            window_id = f"{trace_type}_w{date_token}_{window_spec['slot_token']}"
            source_trace_path = source_dir / build_source_filename(
                trace_type=trace_type, date_text=date_text, slot_token=slot_token
            )
            source_prompt_path = Path(str(source_trace_path).replace(".jsonl", "_prompt.jsonl"))
            windows.append(
                {
                    "window_id": window_id,
                    "trace_type": trace_type,
                    "block_size": TRACE_SPECS[trace_type]["block_size"],
                    "window_start": float(window_spec["window_start"]),
                    "window_end": float(window_spec["window_end"]),
                    "window_index": 6,
                    "sample_order": sample_order,
                    "date": date_text,
                    "slot_token": window_spec["slot_token"],
                    "slot_label": window_spec["slot_label"],
                    "source_trace_path": str(source_trace_path),
                    "source_prompt_path": str(source_prompt_path),
                }
            )
    return windows


def _merge_trace_and_prompt(trace_row: dict[str, Any], prompt_row: dict[str, Any]) -> dict[str, Any]:
    if trace_row.get("chat_id") != prompt_row.get("chat_id") or trace_row.get("turn") != prompt_row.get("turn"):
        raise ValueError(
            "trace/prompt rows are misaligned: "
            f"{trace_row.get('chat_id')}/{trace_row.get('turn')} vs "
            f"{prompt_row.get('chat_id')}/{prompt_row.get('turn')}"
        )
    merged = dict(trace_row)
    prompt = prompt_row.get("prompt")
    if isinstance(prompt, str) and prompt:
        merged["prompt"] = prompt
    return merged


@dataclass
class WindowStats:
    num_requests: int = 0
    sum_input_length: int = 0
    max_input_length: int = 0
    first_request_ts: float | None = None
    last_request_ts: float | None = None
    first_request_index: int | None = None
    last_request_index: int | None = None

    def record(self, row: dict[str, Any]) -> None:
        input_length = int(row.get("input_length") or 0)
        timestamp = float(row["timestamp"])
        if self.num_requests == 0:
            self.first_request_ts = timestamp
            self.first_request_index = 0
        self.last_request_ts = timestamp
        self.last_request_index = self.num_requests
        self.num_requests += 1
        self.sum_input_length += input_length
        self.max_input_length = max(self.max_input_length, input_length)


def materialize_windows(
    windows: list[dict[str, Any]], *, sample_seed: int, traces_dir: Path
) -> dict[str, WindowStats]:
    grouped: dict[tuple[Path, Path], list[dict[str, Any]]] = {}
    for window in windows:
        trace_path = Path(window["source_trace_path"])
        prompt_path = Path(window["source_prompt_path"])
        grouped.setdefault((trace_path, prompt_path), []).append(window)

    stats_by_window = {str(window["window_id"]): WindowStats() for window in windows}
    handles: dict[str, Any] = {}
    final_paths: dict[str, Path] = {}
    temp_paths: dict[str, Path] = {}
    completed = False
    try:
        for window in windows:
            window_id = str(window["window_id"])
            final_path = traces_dir / f"{window_id}.jsonl"
            temp_path = traces_dir / f".{window_id}.jsonl.tmp.{os.getpid()}"
            if temp_path.exists():
                temp_path.unlink()
            final_paths[window_id] = final_path
            temp_paths[window_id] = temp_path
            handles[window_id] = temp_path.open("w", encoding="utf-8")

        for trace_path, prompt_path in sorted(grouped.keys()):
            bucket = grouped[(trace_path, prompt_path)]
            bucket.sort(key=lambda item: (float(item["window_start"]), str(item["window_id"])))
            matched_rows = 0
            root_of: dict[Any, Any] = {}
            with trace_path.open() as trace_handle, prompt_path.open() as prompt_handle:
                for trace_raw, prompt_raw in zip(trace_handle, prompt_handle):
                    trace_raw = trace_raw.strip()
                    if not trace_raw:
                        continue
                    trace_row = json.loads(trace_raw)
                    # Resolve session linkage for every row (even unmatched ones)
                    # so multi-turn chains crossing the window edge still group.
                    session_root = resolve_session_root(trace_row, root_of)
                    timestamp = float(trace_row.get("timestamp") or 0.0)
                    matched_window: dict[str, Any] | None = None
                    for window in bucket:
                        start = float(window["window_start"])
                        end = float(window["window_end"])
                        if start <= timestamp < end:
                            matched_window = window
                            break
                    if matched_window is None:
                        continue
                    prompt_raw = prompt_raw.strip()
                    if not prompt_raw:
                        continue
                    prompt_row = json.loads(prompt_raw)
                    merged = _merge_trace_and_prompt(trace_row, prompt_row)
                    window_id = str(matched_window["window_id"])
                    out = dict(merged)
                    start = float(matched_window["window_start"])
                    out["source_timestamp"] = timestamp
                    out["timestamp"] = timestamp - start
                    out["session_root"] = session_root
                    out["sampling_u"] = session_uniform(
                        seed=sample_seed,
                        window_id=window_id,
                        session_root=session_root,
                    )
                    handles[window_id].write(json.dumps(out, ensure_ascii=False) + "\n")
                    stats_by_window[window_id].record(out)
                    matched_rows += 1
            print(
                f"materialized {trace_path.name} -> matched_rows={matched_rows}",
                flush=True,
            )
        completed = True
    finally:
        for handle in handles.values():
            handle.close()
        if completed:
            for window_id, temp_path in temp_paths.items():
                os.replace(temp_path, final_paths[window_id])
        else:
            for temp_path in temp_paths.values():
                if temp_path.exists():
                    temp_path.unlink()
    return stats_by_window


def build_output_window(
    window: dict[str, Any],
    stats: WindowStats,
    trace_relpath: str,
    *,
    sample_seed: int,
) -> dict[str, Any]:
    output = dict(window)
    output["trace_file"] = trace_relpath
    output["window_start"] = 0.0
    output["window_end"] = float(window["window_end"]) - float(window["window_start"])
    output["num_requests"] = stats.num_requests
    output["sum_input_length"] = stats.sum_input_length
    output["max_input_length"] = stats.max_input_length
    output["num_excluded_too_long"] = 0
    output["sampling_u_field"] = "sampling_u"
    output["sampling_seed"] = int(sample_seed)
    output["sampling_strategy"] = "session_coherent_uniform_score"
    output["first_request_ts"] = stats.first_request_ts
    output["last_request_ts"] = stats.last_request_ts
    output["first_request_index"] = stats.first_request_index
    output["last_request_index"] = stats.last_request_index
    return output


def main() -> int:
    args = parse_args()
    workloads = parse_workloads(args.workloads)
    output_root = args.output_root.resolve()
    if output_root.exists():
        if not args.overwrite:
            raise SystemExit(f"Output root exists: {output_root} (pass --overwrite to replace)")
    output_root.mkdir(parents=True, exist_ok=True)
    traces_dir = output_root / "traces"
    traces_dir.mkdir(parents=True, exist_ok=True)

    source_dirs = {
        "legacy": args.legacy_source.resolve(),
        "thinking": args.thinking_source.resolve(),
    }
    windows = build_windows(source_dirs, workloads=workloads)
    stats_by_window = materialize_windows(
        windows,
        sample_seed=args.sample_seed,
        traces_dir=traces_dir,
    )

    rendered_windows: list[dict[str, Any]] = []
    for window in windows:
        trace_filename = f"{window['window_id']}.jsonl"
        rendered_windows.append(
            build_output_window(
                window,
                stats_by_window[str(window["window_id"])],
                trace_relpath=f"traces/{trace_filename}",
                sample_seed=args.sample_seed,
            )
        )

    windows_payload = {
        "sample_seed": args.sample_seed,
        "source_trace_dirs": {key: str(path) for key, path in source_dirs.items()},
        "u_field": "sampling_u",
        "window_duration_seconds": 600.0,
        "windows": rendered_windows,
    }
    windows_path = output_root / "windows.json"
    windows_tmp_path = output_root / f".windows.json.tmp.{os.getpid()}"
    try:
        windows_tmp_path.write_text(
            json.dumps(windows_payload, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        os.replace(windows_tmp_path, windows_path)
    finally:
        if windows_tmp_path.exists():
            windows_tmp_path.unlink()
    print(output_root)
    print(f"windows={len(rendered_windows)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
