"""Event recorder for AutoDebug probe."""

from __future__ import annotations

import enum
import hashlib
import json
import os
import re
import threading
import time
from dataclasses import dataclass, field
from datetime import datetime, date
from pathlib import Path
from typing import Any, Dict, Optional, Set
from urllib.parse import urlparse

import zstandard as zstd


def _json_default(obj: Any) -> Any:
    """JSON serializer for objects not serializable by default."""
    if isinstance(obj, datetime):
        return obj.isoformat()
    if isinstance(obj, date):
        return obj.isoformat()
    if isinstance(obj, bytes):
        try:
            return obj.decode("utf-8")
        except UnicodeDecodeError:
            import base64
            return {"__bytes__": base64.b64encode(obj).decode("ascii")}
    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")


MAX_RECORD_SIZE_BYTES = 1_048_576  # 1 MB


def _truncate_record(record: Dict[str, Any]) -> tuple[Dict[str, Any], bytes]:
    """Truncate large fields in a record to stay under MAX_RECORD_SIZE_BYTES.

    Returns (record, serialized_bytes) to avoid double-serializing on the hot path.
    """
    serialized = json.dumps(record, separators=(",", ":"), sort_keys=True, default=_json_default)
    encoded = serialized.encode("utf-8")
    if len(encoded) <= MAX_RECORD_SIZE_BYTES:
        return record, encoded

    # Find and truncate the largest string fields
    result = record.copy()
    truncation_marker = "[TRUNCATED]"

    # Fields that are safe to truncate (largest first)
    truncatable_fields = [
        "response_body_bytes", "response_body", "request_body",
        "rows", "result", "fields", "bytes_hex",
        "args", "params_canonical",
    ]

    for field_name in truncatable_fields:
        if field_name in result:
            val = result[field_name]
            if isinstance(val, str) and len(val) > 10000:
                result[field_name] = val[:10000] + truncation_marker
                result["_truncated"] = True
            elif isinstance(val, list) and len(val) > 100:
                result[field_name] = val[:100]
                result["_truncated"] = True
                result["_original_count"] = len(val)

        # Re-check size
        serialized = json.dumps(result, separators=(",", ":"), sort_keys=True, default=_json_default)
        encoded = serialized.encode("utf-8")
        if len(encoded) <= MAX_RECORD_SIZE_BYTES:
            break

    return result, encoded


# Import redaction (handle case where autodebug package not installed)
try:
    from autodebug.redaction import (
        redact_http_record,
        redact_sql_record,
        redact_redis_record,
        redact_grpc_record,
        redact_aws_record,
        redact_mongodb_record,
    )
    REDACTION_ENABLED = True
except ImportError:
    REDACTION_ENABLED = False
    import warnings as _warnings
    _warnings.warn(
        "autodebug.redaction not available — PII redaction is disabled. "
        "Install the autodebug package for production use.",
        stacklevel=2,
    )
    def redact_http_record(record: Dict[str, Any]) -> Dict[str, Any]:
        return record
    def redact_sql_record(record: Dict[str, Any]) -> Dict[str, Any]:
        return record
    def redact_redis_record(record: Dict[str, Any]) -> Dict[str, Any]:
        return record
    def redact_grpc_record(record: Dict[str, Any]) -> Dict[str, Any]:
        return record
    def redact_aws_record(record: Dict[str, Any]) -> Dict[str, Any]:
        return record
    def redact_mongodb_record(record: Dict[str, Any]) -> Dict[str, Any]:
        return record

ORIGINAL_MONOTONIC = time.monotonic


IGNORED_HEADERS_DEFAULT = {
    "date",
    "server",
    "set-cookie",
    "x-request-id",
    "traceparent",
    "tracestate",
}


# ---------------------------------------------------------------------------
# Recording modes (VCR-style)
# ---------------------------------------------------------------------------

class RecordingMode(enum.Enum):
    """Controls when the probe records dependency calls."""
    ALL = "all"
    ONCE = "once"
    NEW_EPISODES = "new_episodes"
    NONE = "none"


def get_recording_mode() -> RecordingMode:
    """Read the recording mode from the AUTODEBUG_RECORD_MODE env var.

    Recognised values (case-insensitive): all, once, new_episodes, new, none.
    Unknown or missing values default to ALL.
    """
    raw = os.environ.get("AUTODEBUG_RECORD_MODE", "").strip().lower()
    if not raw:
        return RecordingMode.ALL
    _mapping = {
        "all": RecordingMode.ALL,
        "once": RecordingMode.ONCE,
        "new_episodes": RecordingMode.NEW_EPISODES,
        "new": RecordingMode.NEW_EPISODES,
        "none": RecordingMode.NONE,
    }
    return _mapping.get(raw, RecordingMode.ALL)


# ---------------------------------------------------------------------------
# SQL normalisation helpers
# ---------------------------------------------------------------------------

# Regex that matches numeric literals (integers and decimals)
_RE_NUMERIC = re.compile(r"\b\d+(?:\.\d+)?\b")
# Regex that matches single-quoted string literals (handles escaped quotes)
_RE_STRING = re.compile(r"'(?:[^'\\]|\\.)*'")


def _normalize_sql(sql: str) -> str:
    """Normalise a SQL query for fingerprinting.

    - Upper-cases the entire string
    - Collapses runs of whitespace to a single space
    - Replaces numeric literals with ``?``
    - Replaces single-quoted string literals with ``'?'``
    """
    sql = sql.upper()
    sql = re.sub(r"\s+", " ", sql).strip()
    sql = _RE_STRING.sub("'?'", sql)
    sql = _RE_NUMERIC.sub("?", sql)
    return sql


# ---------------------------------------------------------------------------
# Fingerprinting helpers
# ---------------------------------------------------------------------------

def _hash16(data: str) -> str:
    """Return the first 16 hex characters of a SHA-256 digest."""
    return hashlib.sha256(data.encode("utf-8")).hexdigest()[:16]


def fingerprint_http(record: Dict[str, Any]) -> str:
    """Create a 16-char fingerprint from an HTTP record's method + path.

    Query parameters are stripped so that ``GET /users?page=1`` and
    ``GET /users?page=2`` produce the same fingerprint.
    """
    method = record.get("method", "").upper()
    url = record.get("url", "")
    parsed = urlparse(url)
    # Use scheme + netloc + path (no query/fragment)
    path = parsed._replace(query="", fragment="").geturl()
    return _hash16(f"{method} {path}")


def fingerprint_sql(record: Dict[str, Any]) -> str:
    """Create a 16-char fingerprint from a normalised SQL query."""
    query = record.get("query", "")
    return _hash16(_normalize_sql(query))


def fingerprint_redis(record: Dict[str, Any]) -> str:
    """Create a 16-char fingerprint from a Redis command + key."""
    command = record.get("command", "").upper()
    key = record.get("key", "")
    return _hash16(f"{command} {key}")


# ---------------------------------------------------------------------------
# Recorded-call index (used by NEW_EPISODES mode)
# ---------------------------------------------------------------------------

class RecordedCallIndex:
    """Tracks which dependency calls have already been recorded."""

    def __init__(self) -> None:
        self.http_fingerprints: Set[str] = set()
        self.sql_fingerprints: Set[str] = set()
        self.redis_fingerprints: Set[str] = set()

    def has_http(self, fp: str) -> bool:
        return fp in self.http_fingerprints

    def has_sql(self, fp: str) -> bool:
        return fp in self.sql_fingerprints

    def has_redis(self, fp: str) -> bool:
        return fp in self.redis_fingerprints


def _load_index_from_dir(run_dir: Path) -> RecordedCallIndex:
    """Build a RecordedCallIndex by reading existing .zst recordings."""
    index = RecordedCallIndex()
    dctx = zstd.ZstdDecompressor()

    http_path = run_dir / "deps_http.jsonl.zst"
    if http_path.exists():
        try:
            with open(http_path, "rb") as f:
                reader = dctx.stream_reader(f)
                content = reader.read().decode("utf-8")
            for line in content.strip().split("\n"):
                if not line:
                    continue
                record = json.loads(line)
                index.http_fingerprints.add(fingerprint_http(record))
        except Exception:
            pass

    sql_path = run_dir / "deps_sql.jsonl.zst"
    if sql_path.exists():
        try:
            with open(sql_path, "rb") as f:
                reader = dctx.stream_reader(f)
                content = reader.read().decode("utf-8")
            for line in content.strip().split("\n"):
                if not line:
                    continue
                record = json.loads(line)
                index.sql_fingerprints.add(fingerprint_sql(record))
        except Exception:
            pass

    redis_path = run_dir / "deps_redis.jsonl.zst"
    if redis_path.exists():
        try:
            with open(redis_path, "rb") as f:
                reader = dctx.stream_reader(f)
                content = reader.read().decode("utf-8")
            for line in content.strip().split("\n"):
                if not line:
                    continue
                record = json.loads(line)
                index.redis_fingerprints.add(fingerprint_redis(record))
        except Exception:
            pass

    return index


# ---------------------------------------------------------------------------
# Stream writer (thread-safe)
# ---------------------------------------------------------------------------

# Intermediate flush interval (seconds) for crash safety.
# The StreamWriter performs a periodic checkpoint flush every this many seconds
# to ensure data is not lost on unexpected process termination (signal, OOM, etc.).
_INTERMEDIATE_FLUSH_INTERVAL = 5.0


@dataclass
class StreamWriter:
    path: Path
    compressor: zstd.ZstdCompressor
    _writer: Optional[zstd.ZstdCompressionWriter] = field(default=None, repr=False)
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
    _write_count: int = field(default=0, repr=False)
    _last_checkpoint_ts: float = field(default=0.0, repr=False)

    def open(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        fh = self.path.open("wb")
        self._writer = self.compressor.stream_writer(fh)
        self._last_checkpoint_ts = time.monotonic()

    def _maybe_intermediate_flush(self) -> None:
        """Perform a periodic intermediate flush for crash safety.

        This checkpoint ensures that if the process is killed (e.g. by a signal
        like SIGKILL or on_crash), we lose at most _INTERMEDIATE_FLUSH_INTERVAL
        seconds of data rather than the entire stream.
        """
        now = time.monotonic()
        if now - self._last_checkpoint_ts >= _INTERMEDIATE_FLUSH_INTERVAL:
            if self._writer is not None:
                self._writer.flush(zstd.FLUSH_BLOCK)
                self._last_checkpoint_ts = now

    def write_jsonl(self, record: Dict[str, Any]) -> None:
        with self._lock:
            if self._writer is None:
                self.open()
            line = json.dumps(record, separators=(",", ":"), sort_keys=True, default=_json_default).encode("utf-8")
            self._writer.write(line + b"\n")
            self._write_count += 1
            self._maybe_intermediate_flush()

    def write_jsonl_raw(self, data: bytes) -> None:
        """Write pre-serialized JSON bytes (avoids double serialization)."""
        with self._lock:
            if self._writer is None:
                self.open()
            self._writer.write(data + b"\n")
            self._write_count += 1
            self._maybe_intermediate_flush()

    @property
    def record_count(self) -> int:
        """Return the number of records written to this stream."""
        return self._write_count

    def close(self) -> None:
        with self._lock:
            if self._writer is None:
                return
            self._writer.flush(zstd.FLUSH_FRAME)
            self._writer.close()
            self._writer = None


# ---------------------------------------------------------------------------
# Recorder
# ---------------------------------------------------------------------------

class Recorder:
    def __init__(self, run_dir: Path, mode: RecordingMode = RecordingMode.ALL) -> None:
        self.run_dir = run_dir
        self.mode = mode
        self._seq = 0
        self._lock = threading.Lock()
        self._closed = False

        # Build call index for NEW_EPISODES mode
        self._index: Optional[RecordedCallIndex] = None
        if mode == RecordingMode.NEW_EPISODES:
            self._index = _load_index_from_dir(run_dir)

        # In ONCE mode, if a previous run completed successfully we suppress all writes.
        # We check for a `.recording_complete` marker rather than just file existence,
        # so that incomplete/crashed runs get re-recorded.
        self._once_skip = False
        if mode == RecordingMode.ONCE:
            if (run_dir / ".recording_complete").exists():
                self._once_skip = True

        # Each stream needs its own compressor to avoid streaming state conflicts
        self.events = StreamWriter(run_dir / "events.jsonl.zst", zstd.ZstdCompressor(level=3))
        self.deps_http = StreamWriter(run_dir / "deps_http.jsonl.zst", zstd.ZstdCompressor(level=3))
        self.deps_sql = StreamWriter(run_dir / "deps_sql.jsonl.zst", zstd.ZstdCompressor(level=3))
        self.deps_redis = StreamWriter(run_dir / "deps_redis.jsonl.zst", zstd.ZstdCompressor(level=3))
        self.deps_grpc = StreamWriter(run_dir / "deps_grpc.jsonl.zst", zstd.ZstdCompressor(level=3))
        self.deps_aws = StreamWriter(run_dir / "deps_aws.jsonl.zst", zstd.ZstdCompressor(level=3))
        self.deps_mongodb = StreamWriter(run_dir / "deps_mongodb.jsonl.zst", zstd.ZstdCompressor(level=3))
        self.time_rng = StreamWriter(run_dir / "time_rng.jsonl.zst", zstd.ZstdCompressor(level=3))
        self.async_task_graph = StreamWriter(run_dir / "async_task_graph.jsonl.zst", zstd.ZstdCompressor(level=3))

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _should_skip(self) -> bool:
        """Return True when recording should be suppressed entirely."""
        return self.mode == RecordingMode.NONE or self._once_skip

    def next_seq(self) -> int:
        with self._lock:
            self._seq += 1
            return self._seq

    def _base_record(self, record_type: str) -> Dict[str, Any]:
        return {
            "seq": self.next_seq(),
            "ts": ORIGINAL_MONOTONIC(),
            "type": record_type,
        }

    # ------------------------------------------------------------------
    # Public query helpers (used by callers for NEW_EPISODES)
    # ------------------------------------------------------------------

    def should_record_http(self, payload: Dict[str, Any]) -> bool:
        """Return True if this HTTP call should be recorded.

        In NEW_EPISODES mode, returns False for previously-seen fingerprints.
        In NONE / ONCE-skip mode, always returns False.
        In ALL mode, always returns True.
        """
        if self._should_skip():
            return False
        if self.mode == RecordingMode.NEW_EPISODES and self._index is not None:
            fp = fingerprint_http(payload)
            return not self._index.has_http(fp)
        return True

    # ------------------------------------------------------------------
    # Recording methods
    # ------------------------------------------------------------------

    def record_event(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record(payload["type"])
        record.update(payload)
        _record, raw = _truncate_record(record)
        self.events.write_jsonl_raw(raw)

    def record_http(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record("HTTP")
        record.update(payload)
        if REDACTION_ENABLED:
            record = redact_http_record(record)
        _record, raw = _truncate_record(record)
        self.deps_http.write_jsonl_raw(raw)

    def record_sql(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record("SQL")
        record.update(payload)
        if REDACTION_ENABLED:
            record = redact_sql_record(record)
        _record, raw = _truncate_record(record)
        self.deps_sql.write_jsonl_raw(raw)

    def record_time_rng(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record(payload["type"])
        record.update(payload)
        _record, raw = _truncate_record(record)
        self.time_rng.write_jsonl_raw(raw)

    def record_redis(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record("REDIS")
        record.update(payload)
        if REDACTION_ENABLED:
            record = redact_redis_record(record)
        _record, raw = _truncate_record(record)
        self.deps_redis.write_jsonl_raw(raw)

    def record_grpc(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record("GRPC")
        record.update(payload)
        if REDACTION_ENABLED:
            record = redact_grpc_record(record)
        _record, raw = _truncate_record(record)
        self.deps_grpc.write_jsonl_raw(raw)

    def record_aws(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record("AWS")
        record.update(payload)
        if REDACTION_ENABLED:
            record = redact_aws_record(record)
        _record, raw = _truncate_record(record)
        self.deps_aws.write_jsonl_raw(raw)

    def record_task_graph(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record(payload["type"])
        record.update(payload)
        _record, raw = _truncate_record(record)
        self.async_task_graph.write_jsonl_raw(raw)

    def record_mongodb(self, payload: Dict[str, Any]) -> None:
        with self._lock:
            if self._closed:
                return
        if self._should_skip():
            return
        record = self._base_record("MONGODB")
        record.update(payload)
        if REDACTION_ENABLED:
            record = redact_mongodb_record(record)
        _record, raw = _truncate_record(record)
        self.deps_mongodb.write_jsonl_raw(raw)

    def close(self) -> None:
        with self._lock:
            if self._closed:
                return
            self._closed = True
        self.events.close()
        self.deps_http.close()
        self.deps_sql.close()
        self.deps_redis.close()
        self.deps_grpc.close()
        self.deps_aws.close()
        self.deps_mongodb.close()
        self.time_rng.close()
        self.async_task_graph.close()

        # Write completion marker for ONCE mode (indicates a clean, non-crashed run)
        if not self._once_skip:
            try:
                (self.run_dir / ".recording_complete").write_text("ok")
            except OSError:
                pass


_recorder: Optional[Recorder] = None


def get_run_dir() -> Path:
    run_dir = os.environ.get("AUTODEBUG_RUN_DIR")
    if run_dir:
        return Path(run_dir)
    run_id = os.environ.get("AUTODEBUG_RUN_ID")
    if not run_id:
        raise RuntimeError("AUTODEBUG_RUN_ID not set")
    return Path(".autodebug") / "runs" / run_id


def get_recorder() -> Recorder:
    global _recorder
    if _recorder is None:
        _recorder = Recorder(get_run_dir())
    return _recorder
