"""Adapter for the Balabit Mouse Dynamics Challenge dataset.

Source: https://github.com/balabit/Mouse-Dynamics-Challenge

Each session is a CSV file with columns:
    record timestamp, client timestamp, button, state, x, y

Where:
    state ∈ {Move, Pressed, Released, Drag, Scroll}
    button ∈ {NoButton, Left, Right, Wheel}

We extract "click-anchored" trajectory segments: each Pressed event
defines a target, and the W ms of Move events preceding it form one
training trace.
"""
from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path

logger = logging.getLogger(__name__)


@dataclass
class MouseEvent:
    """A single mouse event from a Balabit CSV row."""

    t_ms: int        # client timestamp in milliseconds (relative to session start)
    button: str      # "NoButton", "Left", "Right", "Wheel"
    state: str       # "Move", "Pressed", "Released", "Drag", "Scroll"
    x: int
    y: int


@dataclass
class Segment:
    """A click-anchored trajectory segment ready to be written to JSONL."""

    events: list[MouseEvent]   # only Move events, sorted by t_ms ascending
    click_x: int               # the Pressed event's x coordinate
    click_y: int               # the Pressed event's y coordinate
    click_t_ms: int            # the Pressed event's timestamp
    session_id: str            # e.g. "user7_session_42"


def parse_session_csv(path: Path) -> list[MouseEvent]:
    """Parse a Balabit session CSV file into MouseEvent objects.

    Malformed rows are logged and skipped (not raised).
    Client timestamps (seconds, float) are converted to int milliseconds.

    Args:
        path: Path to a Balabit session CSV file.

    Returns:
        List of MouseEvent in original order. Empty list if file is empty.
    """
    import csv as csv_module

    events: list[MouseEvent] = []
    with path.open("r", encoding="utf-8", newline="") as f:
        reader = csv_module.DictReader(f)
        for row_idx, row in enumerate(reader, 2):  # 1-based, header is line 1
            try:
                client_ts = float(row["client timestamp"])
                t_ms = int(round(client_ts * 1000))
                button = row["button"].strip()
                state = row["state"].strip()
                x = int(row["x"])
                y = int(row["y"])
            except (KeyError, ValueError, TypeError) as exc:
                logger.debug("Skipping malformed row %d in %s: %s", row_idx, path.name, exc)
                continue
            events.append(MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y))
    return events


def segment_by_clicks(
    events: list[MouseEvent],
    window_ms: int,
    session_id: str,
) -> list[Segment]:
    """Extract click-anchored segments from a session.

    For each Left-button Pressed event, collect all Move events within
    [click_t - window_ms, click_t) into one segment.

    Args:
        events: Full session events (any state, any order is OK but typically sorted).
        window_ms: How far back to look before each click.
        session_id: String tag attached to every segment for debugging.

    Returns:
        List of Segment, one per Left Pressed event that has at least one preceding Move.
    """
    segments: list[Segment] = []
    for ev in events:
        if ev.state != "Pressed" or ev.button != "Left":
            continue
        click_t = ev.t_ms
        window_start = click_t - window_ms
        moves = [
            m for m in events
            if m.state == "Move" and window_start <= m.t_ms < click_t
        ]
        if not moves:
            continue
        moves.sort(key=lambda m: m.t_ms)
        segments.append(Segment(
            events=moves,
            click_x=ev.x,
            click_y=ev.y,
            click_t_ms=click_t,
            session_id=session_id,
        ))
    return segments


def filter_segments(
    segments: list[Segment],
    min_events: int,
    min_dist: int,
    max_span_ms: int,
    max_gap_ms: int,
    coord_max: int = 5000,
) -> list[Segment]:
    """Drop segments that fail any quality check.

    A segment is dropped if any of these are true:
    - len(events) < min_events
    - Euclidean dist(events[0], (click_x, click_y)) < min_dist
    - events[-1].t_ms - events[0].t_ms > max_span_ms
    - any adjacent Move pair has dt > max_gap_ms (sampling drop-out)
    - any coord (start/end/click) outside [0, coord_max]
    - total arc length < min_dist (high-frequency jitter only)

    Args:
        segments: Candidate segments.
        min_events: Minimum number of Move events.
        min_dist: Minimum start→click pixel distance AND minimum total arc length.
        max_span_ms: Maximum time span of the segment (events[-1] - events[0]).
        max_gap_ms: Maximum allowed gap between adjacent Move events.
        coord_max: Maximum allowed pixel coordinate value (5000 catches multi-monitor anomalies).

    Returns:
        Filtered list, original order preserved.
    """
    import math

    keep: list[Segment] = []
    for seg in segments:
        if len(seg.events) < min_events:
            continue

        sx, sy = seg.events[0].x, seg.events[0].y
        ex, ey = seg.click_x, seg.click_y

        # Coord range check
        if any(c < 0 or c > coord_max for c in (sx, sy, ex, ey)):
            continue

        # Endpoint distance
        dist = math.hypot(ex - sx, ey - sy)
        if dist < min_dist:
            continue

        # Time span
        span = seg.events[-1].t_ms - seg.events[0].t_ms
        if span > max_span_ms:
            continue

        # Gap check + total arc length
        total_arc = 0.0
        bad_gap = False
        for i in range(1, len(seg.events)):
            dt = seg.events[i].t_ms - seg.events[i - 1].t_ms
            if dt > max_gap_ms:
                bad_gap = True
                break
            dx = seg.events[i].x - seg.events[i - 1].x
            dy = seg.events[i].y - seg.events[i - 1].y
            total_arc += math.hypot(dx, dy)
        if bad_gap:
            continue
        if total_arc < min_dist:
            continue

        keep.append(seg)

    return keep


def process_session(
    csv_path: Path,
    output_jsonl: Path,
    config,
) -> int:
    """Convert one Balabit session CSV to JSONL traces, append to output.

    Args:
        csv_path: Path to a Balabit session CSV.
        output_jsonl: Output JSONL file (will be appended to).
        config: BalabitAdapterConfig with window_ms / min_dist / etc.

    Returns:
        Number of valid segments written.
    """
    import math

    session_id = csv_path.stem  # e.g. "session_42"
    events = parse_session_csv(csv_path)
    if not events:
        return 0

    raw_segments = segment_by_clicks(
        events, window_ms=config.window_ms, session_id=session_id
    )
    valid_segments = filter_segments(
        raw_segments,
        min_events=config.min_events,
        min_dist=config.min_dist,
        max_span_ms=config.max_span_ms,
        max_gap_ms=config.max_gap_ms,
    )

    if not valid_segments:
        return 0

    output_jsonl.parent.mkdir(parents=True, exist_ok=True)
    n_written = 0
    with output_jsonl.open("a", encoding="utf-8") as f:
        for seg in valid_segments:
            sx, sy = seg.events[0].x, seg.events[0].y
            ex, ey = seg.click_x, seg.click_y
            dist = math.hypot(ex - sx, ey - sy)
            angle = math.degrees(math.atan2(ey - sy, ex - sx))
            t0 = seg.events[0].t_ms
            record = {
                "meta": {
                    "start": [sx, sy],
                    "end": [ex, ey],
                    "dist": int(round(dist)),
                    "angle": round(angle, 1),
                    "source": "balabit",
                    "session_id": seg.session_id,
                },
                "events": [
                    {"type": "move", "x": e.x, "y": e.y, "t": e.t_ms - t0}
                    for e in seg.events
                ],
            }
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
            n_written += 1
    return n_written


def main(argv: list[str] | None = None) -> int:
    """CLI entry point: convert a directory of Balabit sessions to one JSONL file."""
    import argparse

    from tools.config import BalabitAdapterConfig

    parser = argparse.ArgumentParser(description="Convert Balabit dataset to traces.jsonl format")
    parser.add_argument(
        "--input", type=Path, required=True,
        help="Directory containing Balabit session CSV files (recursive)",
    )
    parser.add_argument(
        "--output", type=Path, default=Path("data/pretrain_traces.jsonl"),
        help="Output JSONL path (default: data/pretrain_traces.jsonl)",
    )
    parser.add_argument("--window-ms", type=int, default=1200)
    parser.add_argument("--min-dist", type=int, default=50)
    parser.add_argument("--min-events", type=int, default=5)
    parser.add_argument("--max-span-ms", type=int, default=5000)
    parser.add_argument("--max-gap-ms", type=int, default=200)
    parser.add_argument(
        "--overwrite", action="store_true",
        help="Truncate output file before writing (default: append)",
    )
    args = parser.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    if not args.input.is_dir():
        logger.error("Input is not a directory: %s", args.input)
        return 2

    config = BalabitAdapterConfig(
        window_ms=args.window_ms,
        min_dist=args.min_dist,
        min_events=args.min_events,
        max_span_ms=args.max_span_ms,
        max_gap_ms=args.max_gap_ms,
    )

    if args.overwrite and args.output.exists():
        args.output.unlink()

    # Recursively walk the input directory looking for session files.
    # Balabit files have no extension; we accept any regular file.
    csv_files = sorted([p for p in args.input.rglob("*") if p.is_file() and not p.name.startswith(".")])
    if not csv_files:
        logger.error("No session files found under %s", args.input)
        return 2

    total = 0
    for i, csv_path in enumerate(csv_files, 1):
        try:
            n = process_session(csv_path, args.output, config)
        except Exception as exc:  # noqa: BLE001
            logger.warning("Skipping %s due to error: %s", csv_path.name, exc)
            continue
        total += n
        if i % 10 == 0 or i == len(csv_files):
            logger.info("Processed %d/%d sessions, %d segments so far", i, len(csv_files), total)

    logger.info("Done. Wrote %d segments to %s", total, args.output)
    return 0


if __name__ == "__main__":
    import sys
    sys.exit(main())