[tab] crowd-pilot/crowd-pilot/serialization_utils.py (python)
#!/usr/bin/env python3
"""
Common utilities for dataset serialization scripts.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional, Tuple, Dict

import difflib
import re
import pandas as pd
from datasets import Dataset, load_dataset


_ANSI_CSI_RE = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]")
_ANSI_OSC_TERMINATED_RE = re.compile(r"\x1b\][\s\S]*?(?:\x07|\x1b\\)")
_ANSI_OSC_LINE_FALLBACK_RE = re.compile(r"\x1b\][^
]*$")
_BRACKETED_PASTE_ENABLE = "\x1b[?2004h"
_BRACKETED_PASTE_DISABLE = "\x1b[?2004l"
_OSC_633 = "\x1b]633;"
_OSC_0 = "\x1b]0;"


@dataclass
class SerializeConfig:
    output_dir: str
    target_chars_per_conversation: int
    target_chars_per_turn: int
    min_session_turns: int
    max_docs: Optional[int]
    csv_root: Optional[str]
    val_ratio: float


@dataclass
class ChunkState:
    """
    Mutable state used while constructing conversation chunks.
    """
    chunks: List[List[Dict[str, str]]]
    max_chars_per_conversation: int
    max_chars_per_turn: int
    current_chunk: List[Dict[str, str]] = field(default_factory=list)
    current_chars: int = 0
    files_opened_in_chunk: set[str] = field(default_factory=set)

    def start_new_chunk(self) -> None:
        if self.current_chunk:
            self.chunks.append(self.current_chunk)
        self.current_chunk = []
        self.current_chars = 0
        self.files_opened_in_chunk.clear()

    def append_turn(self, turn: Dict[str, str]) -> None:
        value = turn.get("value", "")
        # Enforce a per-turn character budget to prevent pathological single
        # turns (e.g., massive file dumps) from creating multi-megabyte chunks.
        if len(value) > self.max_chars_per_turn:
            value = value[: self.max_chars_per_turn]
            turn = dict(turn)
            turn["value"] = value
        turn_len = len(value)
        if (
            self.current_chunk  # only start a new chunk on non-empty current chunk
            and self.current_chars + turn_len > self.max_chars_per_conversation
        ):
            self.start_new_chunk()
        self.current_chunk.append(turn)
        self.current_chars += turn_len

    def maybe_capture_file_contents(
        self,
        file_path: str,
        content: str,
    ) -> None:
        """
        Capture the contents of the given file in the current chunk if it hasn't been opened yet.
        """
        if file_path in self.files_opened_in_chunk:
            return
        cmd = f"cat -n {file_path}"
        self.append_turn({
            "from": "Assistant",
            "value": _fenced_block("bash", _clean_text(cmd)),
        })
        output = _line_numbered_output(content)
        self.append_turn({
            "from": "User",
            "value": f"<stdout>
{output}
</stdout>",
        })
        self.files_opened_in_chunk.add(file_path)


def _clean_text(text: str) -> str:
    # Normalize line endings and strip trailing spaces; preserve tabs/newlines.
    return text.replace("
", "
").replace("
", "
").rstrip()


def _fenced_block(language: Optional[str], content: str) -> str:
    lang = (language or "").lower()
    return f"```{lang}
{content}
```
"


def _apply_change(content: str, offset: int, length: int, new_text: str) -> str:
    # Mirrors crowd_code_player.replay_file.apply_change
    base = str(content)
    text = str(new_text) if pd.notna(new_text) else ""
    text = text.replace("\
", "
").replace("\
", "
")
    if offset > len(base):
        base = base + (" " * (offset - len(base)))
    return base[:offset] + text + base[offset + length:]


def _apply_backspaces(text: str) -> str:
    out: List[str] = []
    for ch in text:
        if ch == "\b":  # \x08
            if out:
                out.pop()
        else:
            out.append(ch)
    return "".join(out)


def _normalize_terminal_output(raw: str) -> str:
    """
    Normalize PTY/terminal output for training:
      - Apply backspaces (\x08)
      - Strip OSC (window title/shell integration) first, keeping BEL/ST terminators intact
      - Resolve carriage returns (
) by keeping the last rewrite per line
      - Strip CSI (coloring etc.)
      - Finally drop any remaining BEL (\x07)
    """
    if not raw:
        return raw
    s = _apply_backspaces(raw)
    # Remove OSC sequences that are properly terminated (BEL or ST)
    s = _ANSI_OSC_TERMINATED_RE.sub("", s)
    # Fallback: drop any unterminated OSC up to end-of-line only
    s = "
".join(_ANSI_OSC_LINE_FALLBACK_RE.sub("", line) for line in s.split("
"))
    # Resolve carriage returns per line:
    # - If there are multiple rewrites, keep the last non-empty chunk
    # - If it's CRLF (ending with '
' before '
'), keep the content before '
'
    resolved_lines: List[str] = []
    for seg in s.split("
"):
        parts = seg.split("
")
        chosen = ""
        # pick last non-empty part if available; else last part
        for p in reversed(parts):
            if p != "":
                chosen = p
                break
        if chosen == "" and parts:
            chosen = parts[-1]
        resolved_lines.append(chosen)
    s = "
".join(resolved_lines)
    # Strip ANSI escape sequences
    s = _ANSI_CSI_RE.sub("", s)
    # Remove any remaining BEL beeps
    s = s.replace("\x07", "")
    return s


def _line_numbered_output(content: str, start_line: Optional[int] = None, end_line: Optional[int] = None) -> str:
    # FIXME (f.srambical): check whether this corresponds **exactly** to the output of cat -n {file_path} | sed -n '{vstart},{vend}p'
    lines = content.splitlines()
    total = len(lines)
    if total == 0:
        return ""
    s = 1 if start_line is None else max(1, min(start_line, total))
    e = total if end_line is None else max(1, min(end_line, total))
    assert e >= s, "End line number cannot be less than start line number! Likely a bug in the line numbering computation."
    buf: List[str] = []
    for idx in range(s, e + 1):
        buf.append(f"{idx:6}	{lines[idx - 1]}")
    return "
".join(buf)


def _compute_viewport(total_lines: int, center_line: int, radius: int) -> Tuple[int, int]:
    if total_lines <= 0:
        return (1, 0)
    start = max(1, center_line - radius)
    end = min(total_lines, center_line + radius)
    assert end >= start, "Viewport cannot have negative width! Likely a bug in the viewport computation."
    return (start, end)


def _escape_single_quotes_for_sed(text: str) -> str:
    # Close quote, add an escaped single quote, reopen quote: '"'"'
    return text.replace("'", "'\"'\"'")


def _compute_changed_block_lines(
    before: str, after: str
) -> Tuple[int, int, int, int, List[str]]:
    """
    Return 1-based start and end line numbers in 'before' that should be
    replaced, 1-based start and end line numbers in 'after' that contain
    the replacement, and the replacement lines from 'after'.

    For pure deletions, the replacement list may be empty.
    """
    before_lines = before.splitlines()
    after_lines = after.splitlines()
    sm = difflib.SequenceMatcher(a=before_lines, b=after_lines, autojunk=False)
    opcodes = [op for op in sm.get_opcodes() if op[0] != "equal"]
    assert opcodes, "Opcode list cannot be empty! Likely a bug in the diff computation."

    first = opcodes[0]
    last = opcodes[-1]
    # i1/i2 refer to 'before' indices, j1/j2 to 'after'
    start_before = max(1, first[1] + 1)
    end_before = last[2]  # no increment since we go from 'exclusive' to 'inclusive' indexing
    start_after = max(1, first[3] + 1)
    end_after = last[4]
    replacement_lines = after_lines[first[3] : last[4]]
    return (start_before, end_before, start_after, end_after, replacement_lines)


def session_to_nemo_conversation_chunks(
    df: pd.DataFrame,
    max_chars_per_conversation: int,
    max_chars_per_turn: int,
    viewport_radius: int = 10,
    normalize_terminal_output: bool = True,
    coalesce_radius: int = 5,
) -> List[List[Dict[str, str]]]:
    """
    Convert a session DataFrame to one or more NeMo conversation chunks.

    - Chunks are created by approximately limiting the total characters
      across all `value` fields to `max_chars_per_conversation`.
    - When a new chunk starts (after the first), the first time a file is
      referenced in that chunk we re-log the full file contents with
      `cat -n <file>` and numbered output so that each chunk is self-contained.
    """
    file_states: Dict[str, str] = {}
    per_file_viewport: Dict[str, Optional[Tuple[int, int]]] = {}

    chunks: List[List[Dict[str, str]]] = []
    chunk_state = ChunkState(
        chunks=chunks,
        max_chars_per_conversation=max_chars_per_conversation,
        max_chars_per_turn=max_chars_per_turn,
    )

    terminal_output_buffer: List[str] = []
    pending_edits_before: Dict[str, Optional[str]] = {}
    pending_edit_regions: Dict[str, Optional[Tuple[int, int]]] = {}

    def _flush_terminal_output_buffer() -> None:
        if not terminal_output_buffer:
            return
        aggregated = "".join(terminal_output_buffer)
        out = aggregated
        if normalize_terminal_output:
            out = _normalize_terminal_output(out)
        cleaned = _clean_text(out)
        if cleaned.strip():
            chunk_state.append_turn({
                "from": "User",
                "value": f"<stdout>
{cleaned}
</stdout>",
            })
        terminal_output_buffer.clear()

    def _flush_pending_edit_for_file(target_file: str) -> None:
        before_snapshot = pending_edits_before.get(target_file)
        if before_snapshot is None:
            return
        after_state = file_states.get(target_file, "")
        if before_snapshot.rstrip("
") == after_state.rstrip("
"):
            pending_edits_before[target_file] = None
            pending_edit_regions[target_file] = None
            return
        (
            start_before,
            end_before,
            start_after,
            end_after,
            repl_lines,
        ) = _compute_changed_block_lines(before_snapshot, after_state)
        before_total_lines = len(before_snapshot.splitlines())
        if end_before < start_before:
            escaped_lines = [_escape_single_quotes_for_sed(line) for line in repl_lines]
            sed_payload = "
".join(escaped_lines)
            if start_before <= max(1, before_total_lines):
                sed_cmd = f"sed -i '{start_before}i\\
{sed_payload}' {target_file}"
            else:
                sed_cmd = f"sed -i '$a\\
{sed_payload}' {target_file}"
        elif not repl_lines:
            sed_cmd = f"sed -i '{start_before},{end_before}d' {target_file}"
        else:
            escaped_lines = [_escape_single_quotes_for_sed(line) for line in repl_lines]
            sed_payload = "
".join(escaped_lines)
            sed_cmd = f"sed -i '{start_before},{end_before}c\\
{sed_payload}' {target_file}"
        total_lines = len(after_state.splitlines())
        center = (start_after + end_after) // 2
        vp = _compute_viewport(total_lines, center, viewport_radius)
        per_file_viewport[target_file] = vp
        vstart, vend = vp
        chunk_state.maybe_capture_file_contents(target_file, before_snapshot)
        chained_cmd = f"{sed_cmd} && cat -n {target_file} | sed -n '{vstart},{vend}p'"
        chunk_state.append_turn({
            "from": "Assistant",
            "value": _fenced_block("bash", _clean_text(chained_cmd)),
        })
        viewport_output = _line_numbered_output(after_state, vstart, vend)
        chunk_state.append_turn({
            "from": "User",
            "value": f"<stdout>
{viewport_output}
</stdout>",
        })
        pending_edits_before[target_file] = None
        pending_edit_regions[target_file] = None

    def _flush_all_pending_edits() -> None:
        for fname in list(pending_edits_before.keys()):
            _flush_pending_edit_for_file(fname)

    for i in range(len(df)):
        row = df.iloc[i]
        file_path: str = row["File"]
        event_type = row["Type"]

        match event_type:
            case "tab":
                _flush_all_pending_edits()
                _flush_terminal_output_buffer()
                text = row["Text"]
                if pd.notna(text):
                    content = str(text).replace("\
", "
").replace("\
", "
")
                    file_states[file_path] = content
                    cmd = f"cat -n {file_path}"
                    chunk_state.append_turn({
                        "from": "Assistant",
                        "value": _fenced_block("bash", _clean_text(cmd)),
                    })
                    output = _line_numbered_output(content)
                    chunk_state.append_turn({
                        "from": "User",
                        "value": f"<stdout>
{output}
</stdout>",
                    })
                    chunk_state.files_opened_in_chunk.add(file_path)
                else:
                    # File switch without content snapshot: show current viewport only
                    content = file_states.get(file_path, "")
                    total_lines = len(content.splitlines())
                    vp = per_file_viewport.get(file_path)
                    if not vp or vp[1] == 0:
                        vp = _compute_viewport(total_lines, 1, viewport_radius)
                        per_file_viewport[file_path] = vp
                    if vp:
                        vstart, vend = vp
                        chunk_state.maybe_capture_file_contents(file_path, content)
                        cmd = f"cat -n {file_path} | sed -n '{vstart},{vend}p'"
                        chunk_state.append_turn({
                            "from": "Assistant",
                            "value": _fenced_block("bash", _clean_text(cmd)),
                        })
                        viewport_output = _line_numbered_output(content, vstart, vend)
                        chunk_state.append_turn({
                            "from": "User",
                            "value": f"<stdout>
{viewport_output}
</stdout>",
                        })

            case "content":
                _flush_terminal_output_buffer()
                offset = int(row["RangeOffset"])
                length = int(row["RangeLength"])
                new_text = row["Text"]
                before = file_states.get(file_path, "")
                # Approximate current edit region in line space
                new_text_str = str(new_text) if pd.notna(new_text) else ""
                start_line_current = before[:offset].count("
") + 1
                deleted_chunk = before[offset:offset + length]
                lines_added = new_text_str.count("
")
                lines_deleted = deleted_chunk.count("
")
                region_start = start_line_current
                region_end = start_line_current + max(lines_added, lines_deleted, 0)
                # Flush pending edits if this edit is far from the pending region
                current_region = pending_edit_regions.get(file_path)
                if current_region is not None:
                    rstart, rend = current_region
                    if region_start < (rstart - coalesce_radius) or region_start > (rend + coalesce_radius):
                        _flush_pending_edit_for_file(file_path)
                        current_region = None
                after = _apply_change(before, offset, length, new_text)
                if pending_edits_before.get(file_path) is None:
                    pending_edits_before[file_path] = before
                # Update/initialize region union
                if current_region is None:
                    pending_edit_regions[file_path] = (region_start, max(region_start, region_end))
                else:
                    rstart, rend = current_region
                    pending_edit_regions[file_path] = (min(rstart, region_start), max(rend, region_end))
                file_states[file_path] = after

            case "selection_command" | "selection_mouse" | "selection_keyboard":
                # During an edit burst (pending edits), suppress flush and viewport emissions
                if pending_edits_before.get(file_path) is None:
                    _flush_terminal_output_buffer()
                else:
                    # Skip emitting viewport while edits are pending to avoid per-keystroke sed/cat spam
                    continue
                offset = int(row["RangeOffset"])
                content = file_states.get(file_path, "")
                total_lines = len(content.splitlines())
                target_line = content[:offset].count("
") + 1
                vp = per_file_viewport.get(file_path)
                should_emit = False
                if not vp or vp[1] == 0:
                    vp = _compute_viewport(total_lines, target_line, viewport_radius)
                    per_file_viewport[file_path] = vp
                    should_emit = True
                else:
                    vstart, vend = vp
                    if target_line < vstart or target_line > vend:
                        vp = _compute_viewport(total_lines, target_line, viewport_radius)
                        per_file_viewport[file_path] = vp
                        should_emit = True
                if should_emit and vp:
                    vstart, vend = vp
                    chunk_state.maybe_capture_file_contents(file_path, content)
                    cmd = f"cat -n {file_path} | sed -n '{vstart},{vend}p'"
                    chunk_state.append_turn({
                        "from": "Assistant",
                        "value": _fenced_block("bash", _clean_text(cmd)),
                    })
                    viewport_output = _line_numbered_output(content, vstart, vend)
                    chunk_state.append_turn({
                        "from": "User",
                        "value": f"<stdout>
{viewport_output}
</stdout>",
                    })

            case "terminal_command":
                _flush_all_pending_edits()
                _flush_terminal_output_buffer()
                command = row["Text"]
                command_str = str(command).replace("\
", "
").replace("\
", "
")
                chunk_state.append_turn({
                    "from": "Assistant",
                    "value": _fenced_block("bash", _clean_text(command_str)),
                })

            case "terminal_output":
                output = row["Text"]
                raw_output = str(output).replace("\
", "
").replace("\
", "
")
                terminal_output_buffer.append(raw_output)

            case "terminal_focus":
                _flush_all_pending_edits()
                _flush_terminal_output_buffer()
                # No-op for bash transcript; focus changes don't emit commands/output
                pass

            case "git_branch_checkout":
                _flush_all_pending_edits()
                _flush_terminal_output_buffer()
                branch_info = row["Text"]
                branch_str = str(branch_info).replace("\
", "
").replace("\
", "
")
                cleaned = _clean_text(branch_str)
                m = re.search(r"to '([^']+)'", cleaned)
                if not m:
                    raise ValueError(f"Could not extract branch name from git checkout message: {cleaned}")
                branch_name = m.group(1).strip()
                # Safe-quote branch if it contains special characters
                if re.search(r"[^A-Za-z0-9._/\\-]", branch_name):
                    branch_name = "'" + branch_name.replace("'", "'\"'\"'") + "'"
                cmd = f"git checkout {branch_name}"
                chunk_state.append_turn({
                    "from": "Assistant",
                    "value": _fenced_block("bash", _clean_text(cmd)),
                })

            case _:
                raise ValueError(f"Unknown event type: {event_type}")

    _flush_all_pending_edits()
    _flush_terminal_output_buffer()
    if chunk_state.current_chunk:
        chunks.append(chunk_state.current_chunk)
    return chunks



def load_hf_csv(hf_path: str, split: str) -> Dataset:
    loaded = load_dataset(hf_path, split=split)

    assert isinstance(loaded, Dataset), "Expected a Dataset from load_dataset"
    return loaded


def _discover_local_sessions(root: Path) -> List[Path]:
    # Recursively find all CSV files
    paths: List[Path] = []
    for p in root.rglob("*.csv"):
        if p.is_file():
            paths.append(p)
    paths.sort()
    return paths