"""Transcript diff tool for comparing two runs.

This module compares two transcripts to identify where they diverge,
helping debug non-deterministic tests or unexpected behavior changes.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

try:
    from api.metrics import inc as _metrics_inc
except Exception:  # pragma: no cover
    _metrics_inc = None


def _inc_metric(name: str, value: int = 1) -> None:
    if _metrics_inc is None:
        return
    try:
        _metrics_inc(name, value)
    except Exception:
        pass


@dataclass
class DiffEntry:
    """A single difference between two transcripts."""

    seq: int
    record_type: str
    field: str
    left_value: Any
    right_value: Any
    is_first_divergence: bool = False

    def __str__(self) -> str:
        marker = " ^ FIRST DIVERGENCE" if self.is_first_divergence else ""
        return (
            f"Seq {self.seq}: {self.record_type} - {self.field}\n"
            f"  - left:  {self._truncate(self.left_value)}\n"
            f"  + right: {self._truncate(self.right_value)}{marker}"
        )

    def _truncate(self, value: Any, max_len: int = 80) -> str:
        s = str(value)
        if len(s) > max_len:
            return s[: max_len - 3] + "..."
        return s


@dataclass
class DiffResult:
    """Result of comparing two transcripts."""

    left_run_id: str
    right_run_id: str
    left_count: int
    right_count: int
    differences: List[DiffEntry]
    identical_count: int
    first_divergence_seq: Optional[int]

    @property
    def has_differences(self) -> bool:
        return len(self.differences) > 0


_UUID_RE = re.compile(
    r"^[a-f0-9]{8}-[a-f0-9]{4}-[1-5][a-f0-9]{3}-[89ab][a-f0-9]{3}-[a-f0-9]{12}$",
    re.IGNORECASE,
)
_ISO_TS_RE = re.compile(
    r"^\d{4}-\d{2}-\d{2}[Tt ]\d{2}:\d{2}:\d{2}(\.\d+)?([Zz]|[+\-]\d{2}:\d{2})?$"
)
_VOLATILE_FIELDS = {
    "ts",
    "timestamp",
    "duration",
    "request_id",
    "trace_id",
    "correlation_id",
    "span_id",
    "created_at",
    "updated_at",
}
_HTTP_PAYLOAD_FIELDS = {
    "body",
    "json",
    "payload",
    "data",
    "response",
    "response_body",
    "request_body",
}
_SQL_FIELDS = {"query", "sql", "sql_canonical"}


def _looks_volatile_key(key: str) -> bool:
    k = key.lower()
    if k in _VOLATILE_FIELDS:
        return True
    return (
        "timestamp" in k
        or k.endswith("_id")
        or k == "id"
        or "uuid" in k
        or k.endswith("_at")
    )


def _canonicalize_scalar(value: Any) -> Any:
    if isinstance(value, str):
        stripped = value.strip()
        if _UUID_RE.match(stripped):
            return "<UUID>"
        if _ISO_TS_RE.match(stripped):
            return "<TIMESTAMP>"
    return value


def _canonicalize_json(value: Any) -> Any:
    if isinstance(value, dict):
        out: Dict[str, Any] = {}
        for key in sorted(value.keys()):
            if _looks_volatile_key(key):
                continue
            out[key] = _canonicalize_json(value[key])
        return out
    if isinstance(value, list):
        items = [_canonicalize_json(v) for v in value]
        # Treat payload arrays as order-insensitive to reduce false divergences.
        # Cap sorting to avoid O(n log n * serialization) on very large arrays.
        if len(items) <= 500:
            return sorted(
                items,
                key=lambda x: json.dumps(x, sort_keys=True, default=str),
            )
        return items
    return _canonicalize_scalar(value)


def _normalize_sql_query(sql: str) -> str:
    q = sql.strip().rstrip(";")
    if not q:
        return q
    # Normalize literals and whitespace (handle escaped quotes like 'it''s').
    q = re.sub(r"'(?:[^']|'')*'", "'?'", q)
    q = re.sub(r'"(?:[^"]|"")*"', '"?"', q)
    q = re.sub(
        r"\b[a-f0-9]{8}-[a-f0-9]{4}-[1-8][a-f0-9]{3}-[89ab][a-f0-9]{3}-[a-f0-9]{12}\b",
        "?",
        q,
        flags=re.IGNORECASE,
    )
    q = re.sub(r"\b\d+(\.\d+)?\b", "?", q)
    q = re.sub(r"\s+", " ", q).strip().lower()

    # Canonicalize simple SELECT projection order.
    m = re.match(r"^select\s+(.+?)\s+from\s+(.+)$", q, flags=re.IGNORECASE)
    if m:
        projection = m.group(1).strip()
        tail = m.group(2).strip()
        if projection != "*" and "(" not in projection and ")" not in projection:
            cols = [c.strip() for c in projection.split(",") if c.strip()]
            if len(cols) > 1:
                projection = ", ".join(sorted(cols))
                q = f"select {projection} from {tail}"

    return q[:500]


def _normalize_record(record: Dict[str, Any]) -> Dict[str, Any]:
    """Normalize a record for comparison by removing volatile fields."""
    normalized: Dict[str, Any] = {}
    record_type = str(record.get("type", "UNKNOWN")).upper()

    for key, value in record.items():
        if _looks_volatile_key(key):
            continue
        # Normalize type field to uppercase so "sql" vs "SQL" doesn't cause false diffs
        if key == "type":
            normalized[key] = record_type
            continue
        if record_type == "SQL" and key.lower() in _SQL_FIELDS and isinstance(value, str):
            normalized[key] = _normalize_sql_query(value)
            continue
        if record_type == "HTTP" and key.lower() in _HTTP_PAYLOAD_FIELDS:
            normalized[key] = _canonicalize_json(value)
            continue
        if isinstance(value, (dict, list)):
            normalized[key] = _canonicalize_json(value)
            continue
        normalized[key] = _canonicalize_scalar(value)

    return normalized


def _compare_records(
    left: Dict[str, Any],
    right: Dict[str, Any],
    seq: int,
) -> List[DiffEntry]:
    """Compare two records and return list of differences."""
    differences = []
    record_type = left.get("type", right.get("type", "UNKNOWN"))

    left_norm = _normalize_record(left)
    right_norm = _normalize_record(right)

    all_keys = set(left_norm.keys()) | set(right_norm.keys())

    for key in sorted(all_keys):
        left_val = left_norm.get(key)
        right_val = right_norm.get(key)

        if left_val != right_val:
            differences.append(
                DiffEntry(
                    seq=seq,
                    record_type=record_type,
                    field=key,
                    left_value=left_val,
                    right_value=right_val,
                )
            )

    return differences


def diff_transcripts(
    left: List[Dict[str, Any]],
    right: List[Dict[str, Any]],
    left_run_id: str = "left",
    right_run_id: str = "right",
) -> DiffResult:
    """Compare two transcripts and identify differences.

    Args:
        left: First transcript (list of records)
        right: Second transcript (list of records)
        left_run_id: Identifier for first transcript
        right_run_id: Identifier for second transcript

    Returns:
        DiffResult with all differences found
    """
    _inc_metric("replay.diff.runs")

    # Index by sequence number, preserving duplicates
    left_by_seq: Dict[Any, List[Dict[str, Any]]] = {}
    for i, r in enumerate(left):
        left_by_seq.setdefault(r.get("seq", i), []).append(r)
    right_by_seq: Dict[Any, List[Dict[str, Any]]] = {}
    for i, r in enumerate(right):
        right_by_seq.setdefault(r.get("seq", i), []).append(r)

    all_seqs = sorted(set(left_by_seq.keys()) | set(right_by_seq.keys()))

    differences: List[DiffEntry] = []
    identical_count = 0
    first_divergence_seq = None

    for seq in all_seqs:
        left_recs = left_by_seq.get(seq, [])
        right_recs = right_by_seq.get(seq, [])

        paired = max(len(left_recs), len(right_recs))
        for idx in range(paired):
            left_rec = left_recs[idx] if idx < len(left_recs) else None
            right_rec = right_recs[idx] if idx < len(right_recs) else None

            if left_rec is None:
                # Record only in right
                differences.append(
                    DiffEntry(
                        seq=seq,
                        record_type=right_rec.get("type", "UNKNOWN"),
                        field="<record>",
                        left_value="<missing>",
                        right_value=right_rec,
                    )
                )
                if first_divergence_seq is None:
                    first_divergence_seq = seq
                    differences[-1].is_first_divergence = True

            elif right_rec is None:
                # Record only in left
                differences.append(
                    DiffEntry(
                        seq=seq,
                        record_type=left_rec.get("type", "UNKNOWN"),
                        field="<record>",
                        left_value=left_rec,
                        right_value="<missing>",
                    )
                )
                if first_divergence_seq is None:
                    first_divergence_seq = seq
                    differences[-1].is_first_divergence = True

            else:
                # Both have record, compare fields
                record_diffs = _compare_records(left_rec, right_rec, seq)
                if record_diffs:
                    if first_divergence_seq is None:
                        first_divergence_seq = seq
                        record_diffs[0].is_first_divergence = True
                    differences.extend(record_diffs)
                else:
                    identical_count += 1

    result = DiffResult(
        left_run_id=left_run_id,
        right_run_id=right_run_id,
        left_count=len(left),
        right_count=len(right),
        differences=differences,
        identical_count=identical_count,
        first_divergence_seq=first_divergence_seq,
    )
    if result.has_differences:
        _inc_metric("replay.diff.divergent")
        if result.differences:
            first_type = str(result.differences[0].record_type or "unknown").lower()
            _inc_metric(f"replay.diff.first_type.{first_type}")
    else:
        _inc_metric("replay.diff.identical")
    return result


def diff_by_type(
    left: List[Dict[str, Any]],
    right: List[Dict[str, Any]],
) -> Dict[str, DiffResult]:
    """Compare transcripts grouped by record type.

    This is useful when you want to see HTTP differences separately
    from SQL differences, etc.

    Args:
        left: First transcript
        right: Second transcript

    Returns:
        Dict mapping record type to DiffResult
    """
    # Group by type
    left_by_type: Dict[str, List[Dict[str, Any]]] = {}
    right_by_type: Dict[str, List[Dict[str, Any]]] = {}

    for record in left:
        rtype = record.get("type", "UNKNOWN")
        left_by_type.setdefault(rtype, []).append(record)

    for record in right:
        rtype = record.get("type", "UNKNOWN")
        right_by_type.setdefault(rtype, []).append(record)

    all_types = set(left_by_type.keys()) | set(right_by_type.keys())

    results = {}
    for rtype in sorted(all_types):
        left_records = left_by_type.get(rtype, [])
        right_records = right_by_type.get(rtype, [])
        results[rtype] = diff_transcripts(
            left_records, right_records, f"left:{rtype}", f"right:{rtype}"
        )

    return results


def load_transcript(path: Path) -> List[Dict[str, Any]]:
    """Load transcript from a JSONL or JSONL.zst file.

    Malformed lines are skipped and counted rather than aborting the load.
    """
    import logging

    _log = logging.getLogger(__name__)
    records: List[Dict[str, Any]] = []
    bad_lines = 0

    def _parse_line(line: str, line_num: int) -> None:
        nonlocal bad_lines
        stripped = line.strip()
        if not stripped:
            return
        try:
            records.append(json.loads(stripped))
        except (json.JSONDecodeError, ValueError) as exc:
            bad_lines += 1
            if bad_lines <= 5:
                _log.warning("Skipping bad JSONL record at %s:%d: %s", path, line_num, exc)

    if path.suffix == ".zst" or path.name.endswith(".jsonl.zst"):
        try:
            import zstandard as zstd
        except ImportError:
            raise RuntimeError(
                "zstandard package required to read .zst files. Install with: pip install zstandard"
            )
        with path.open("rb") as fh:
            dctx = zstd.ZstdDecompressor()
            with dctx.stream_reader(fh) as reader:
                import io
                text_reader = io.TextIOWrapper(reader, encoding="utf-8")
                for line_num, line in enumerate(text_reader, 1):
                    _parse_line(line, line_num)
    else:
        with path.open("r", encoding="utf-8") as f:
            for line_num, line in enumerate(f, 1):
                _parse_line(line, line_num)

    if bad_lines:
        _log.warning("Loaded %s: %d records OK, %d bad lines skipped", path, len(records), bad_lines)
        _inc_metric("transcript.load.bad_lines", bad_lines)
    return records


def format_diff_report(result: DiffResult, max_diffs: int = 20) -> str:
    """Format diff result as a human-readable report.

    Args:
        result: DiffResult to format
        max_diffs: Maximum number of differences to show

    Returns:
        Formatted report string
    """
    lines = [
        "",
        f"Transcript Diff: {result.left_run_id} vs {result.right_run_id}",
        "=" * 60,
        "",
        f"Left records:  {result.left_count}",
        f"Right records: {result.right_count}",
        f"Identical:     {result.identical_count}",
        f"Differences:   {len(result.differences)}",
        "",
    ]

    if result.first_divergence_seq is not None:
        lines.append(f"First divergence at seq {result.first_divergence_seq}")
        lines.append("")

    if not result.has_differences:
        lines.append("Transcripts are identical.")
        return "\n".join(lines)

    lines.append("Differences:")
    lines.append("-" * 40)

    shown = 0
    for diff in result.differences:
        if shown >= max_diffs:
            remaining = len(result.differences) - shown
            lines.append(f"\n... and {remaining} more differences")
            break

        lines.append(str(diff))
        lines.append("")
        shown += 1

    # Summary by type
    by_type: Dict[str, int] = {}
    for diff in result.differences:
        by_type[diff.record_type] = by_type.get(diff.record_type, 0) + 1

    lines.append("")
    lines.append("Summary by type:")
    for rtype, count in sorted(by_type.items(), key=lambda x: -x[1]):
        lines.append(f"  {rtype}: {count} differences")

    return "\n".join(lines)


def format_type_diff_report(results: Dict[str, DiffResult]) -> str:
    """Format type-grouped diff results."""
    lines = [
        "",
        "Transcript Diff by Type",
        "=" * 60,
        "",
    ]

    for rtype, result in results.items():
        if not result.has_differences:
            lines.append(f"{rtype}: identical ({result.identical_count} records)")
        else:
            lines.append(
                f"{rtype}: {len(result.differences)} differences "
                f"(first at seq {result.first_divergence_seq})"
            )

    lines.append("")

    # Show details for types with differences
    for rtype, result in results.items():
        if result.has_differences:
            lines.append(f"\n{rtype} Details:")
            lines.append("-" * 40)
            for diff in result.differences[:5]:
                lines.append(str(diff))
                lines.append("")
            if len(result.differences) > 5:
                lines.append(f"... and {len(result.differences) - 5} more")

    return "\n".join(lines)
