"""Transcript loading and sequence enforcement."""

from __future__ import annotations

import json
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Set, Tuple

try:
    import zstandard as zstd
    HAS_ZSTD = True
except ImportError:
    HAS_ZSTD = False

STREAMING_THRESHOLD_BYTES = 10 * 1024 * 1024


_DIVERGENCE_HINTS = {
    "MISMATCH": "The test made a different I/O call than expected. Check if test behavior depends on time, random values, or environment variables.",
    "MISSING_DEP": "The test made more {type} calls than were recorded. The test may have conditional logic that behaves differently across runs.",
    "EXTRA_CALL": "The test made an I/O call that wasn't in the recording.",
    "WRONG_ORDER": "I/O calls happened in a different sequence than recorded.",
    "PAYLOAD_MISMATCH": "Same call target, different request body. Look for timestamps, UUIDs, or nonces in the request.",
    "NETWORK_ACCESS": "Real network access attempted during replay. Ensure all HTTP/gRPC/DB calls are captured in the transcript.",
    "SQL_DIVERGENCE": "SQL query ordering diverged. Check for non-deterministic query ordering or transaction boundaries.",
    "SQL_MISMATCH": "SQL query text doesn't match the recording. The test may be generating dynamic SQL.",
    "ASYNC_ORDER": "Async task completion order diverged from capture. Tasks completed in a different sequence, likely due to non-deterministic scheduling in asyncio.gather() or create_task().",
    "UNCAPTURED_ENV_ACCESS": "Environment variable '{var}' was accessed but not captured during recording.",
    "UNCAPTURED_FILE_READ": "File '{path}' was read but not captured during recording.",
}


class ReplayDivergence(RuntimeError):
    def __init__(self, code: str, detail: Dict[str, Any]):
        self.code = code
        self.detail = detail
        super().__init__(self._format())

    def _format(self) -> str:
        lines = [f"ReplayDivergence [{self.code}]"]
        for k, v in self.detail.items():
            v_str = str(v)
            if len(v_str) > 200:
                v_str = v_str[:197] + "..."
            lines.append(f"  {k}: {v_str}")
        hint = _DIVERGENCE_HINTS.get(self.code)
        if hint:
            try:
                hint = hint.format(**self.detail)
            except (KeyError, IndexError):
                pass
            lines.append(f"  Hint: {hint}")
        return "\n".join(lines)


@dataclass
class TranscriptIndex:
    """Type-specific transcript index with optional task-aware matching.

    Each record type (HTTP, SQL, TIME, RNG) has its own position counter
    to handle framework differences (e.g., pytest version making different
    timing calls) without affecting order of critical dependencies (HTTP, SQL).

    When a task graph is loaded (v2.0 bundles), records can also be matched
    by task_id for deterministic async replay. Falls back to type-only
    matching for v1.0 bundles.

    Records are pre-indexed by type at construction for O(1) lookup.
    """
    records: List[Dict[str, Any]]
    positions: Dict[str, int] = field(default_factory=dict)
    _type_index: Dict[str, List[Dict[str, Any]]] = field(default_factory=dict, init=False, repr=False)
    _lock: threading.RLock = field(default_factory=threading.RLock, init=False, repr=False)
    # Task graph for async-aware replay
    task_graph: List[Dict[str, Any]] = field(default_factory=list)
    _task_type_index: Dict[str, List[Dict[str, Any]]] = field(default_factory=dict, init=False, repr=False)
    _task_type_positions: Dict[str, int] = field(default_factory=dict, init=False, repr=False)

    def __post_init__(self):
        with self._lock:
            self._type_index = {}
            for r in self.records:
                rtype = r.get("type")
                if rtype:
                    self._type_index.setdefault(rtype, []).append(r)
            # Build task-aware index: key = "task_id:type"
            self._task_type_index = {}
            self._task_type_positions = {}
            for r in self.records:
                rtype = r.get("type")
                task_id = r.get("task_id")
                if rtype and task_id is not None:
                    key = f"{task_id}:{rtype}"
                    self._task_type_index.setdefault(key, []).append(r)

    @property
    def has_task_graph(self) -> bool:
        """Whether this index has async task graph data."""
        return len(self.task_graph) > 0

    def _records_of_type(self, record_type: str) -> List[Dict[str, Any]]:
        return self._type_index.get(record_type, [])

    def expect(self, record_type: str, task_id: Optional[int] = None) -> Dict[str, Any]:
        """Return the next expected record of the given type.

        If task_id is provided and task graph data exists, matches by
        task_id + type for deterministic async ordering. Falls back to
        type-only matching when no task graph is present.
        """
        with self._lock:
            # Task-aware matching when task graph exists
            if task_id is not None and self.has_task_graph:
                key = f"{task_id}:{record_type}"
                pos = self._task_type_positions.get(key, 0)
                task_records = self._task_type_index.get(key, [])
                if pos < len(task_records):
                    self._task_type_positions[key] = pos + 1
                    return task_records[pos]
                # Fall through to type-based if task-specific not found

            # Type-based matching (v1 compat / fallback)
            pos = self.positions.get(record_type, 0)
            type_records = self._records_of_type(record_type)
            if pos >= len(type_records):
                raise ReplayDivergence("MISSING_DEP", {"type": record_type, "pos": pos})
            record = type_records[pos]
            self.positions[record_type] = pos + 1
            return record

    def has_more(self, record_type: str) -> bool:
        """Check if more records of this type are available."""
        with self._lock:
            pos = self.positions.get(record_type, 0)
            return pos < len(self._records_of_type(record_type))

    def get_task_completion_order(self) -> List[int]:
        """Return task IDs in the order they completed during capture."""
        completions = [
            r for r in self.task_graph
            if r.get("type") == "TASK_COMPLETE"
        ]
        completions.sort(key=lambda r: r.get("seq", 0))
        return [r["task_id"] for r in completions]


def _stream_jsonl_file(path: Path) -> Iterator[Dict[str, Any]]:
    """Stream JSONL records from a file, handling zstd compression if present."""
    zst_path = Path(str(path) + ".zst")
    if zst_path.exists():
        if not HAS_ZSTD:
            raise RuntimeError(f"zstandard package required to read {zst_path}")
        with open(zst_path, "rb") as f:
            dctx = zstd.ZstdDecompressor()
            with dctx.stream_reader(f) as reader:
                data = reader.read()
            for line in data.decode("utf-8").splitlines():
                line = line.strip()
                if line:
                    yield json.loads(line)
        return

    if not path.exists():
        return

    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                yield json.loads(line)


def _read_jsonl_file(path: Path) -> List[Dict[str, Any]]:
    """Read a JSONL file, handling zstd compression if present."""
    records: List[Dict[str, Any]] = []

    # Check for zstd compressed version first
    zst_path = Path(str(path) + ".zst")
    if zst_path.exists():
        if not HAS_ZSTD:
            raise RuntimeError(f"zstandard package required to read {zst_path}")
        with open(zst_path, "rb") as f:
            dctx = zstd.ZstdDecompressor()
            # Use stream reader for streaming compressed data
            with dctx.stream_reader(f) as reader:
                data = reader.read()
            for line in data.decode("utf-8").splitlines():
                line = line.strip()
                if line:
                    records.append(json.loads(line))
        return records

    # Fall back to plain jsonl
    if path.exists():
        with path.open("r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if line:
                    records.append(json.loads(line))
    return records


def _estimate_transcript_size(transcripts_dir: Path) -> int:
    """Estimate total transcript size from jsonl or zst files."""
    total = 0
    for name in [
        "events.jsonl",
        "deps_http.jsonl",
        "deps_sql.jsonl",
        "deps_redis.jsonl",
        "deps_grpc.jsonl",
        "deps_aws.jsonl",
        "deps_mongodb.jsonl",
        "time_rng.jsonl",
        "async_task_graph.jsonl",
    ]:
        path = transcripts_dir / name
        zst_path = Path(str(path) + ".zst")
        if zst_path.exists():
            total += zst_path.stat().st_size
        elif path.exists():
            total += path.stat().st_size
    return total


class StreamingTranscriptIndex:
    """Transcript index that loads records lazily by type."""

    _type_to_files: Dict[str, Tuple[str, ...]] = {
        "EVENT": ("events.jsonl",),
        "HTTP": ("deps_http.jsonl",),
        "SQL": ("deps_sql.jsonl",),
        "REDIS": ("deps_redis.jsonl",),
        "GRPC": ("deps_grpc.jsonl",),
        "AWS": ("deps_aws.jsonl",),
        "MONGODB": ("deps_mongodb.jsonl",),
        "TIME": ("time_rng.jsonl",),
        "RNG": ("time_rng.jsonl",),
    }

    def __init__(self, transcripts_dir: Path):
        self.transcripts_dir = transcripts_dir
        self.positions: Dict[str, int] = {}
        self._type_index: Dict[str, List[Dict[str, Any]]] = {}
        self._loaded_types: Set[str] = set()
        self._lock = threading.RLock()
        # Load task graph eagerly (small file)
        self.task_graph: List[Dict[str, Any]] = list(
            _stream_jsonl_file(transcripts_dir / "async_task_graph.jsonl")
        )

    @property
    def has_task_graph(self) -> bool:
        return len(self.task_graph) > 0

    def get_task_completion_order(self) -> List[int]:
        completions = [
            r for r in self.task_graph
            if r.get("type") == "TASK_COMPLETE"
        ]
        completions.sort(key=lambda r: r.get("seq", 0))
        return [r["task_id"] for r in completions]

    def _load_type(self, record_type: str) -> None:
        with self._lock:
            if record_type in self._loaded_types:
                return
            files = self._type_to_files.get(record_type)
            if not files:
                self._type_index[record_type] = []
                self._loaded_types.add(record_type)
                return
            records: List[Dict[str, Any]] = []
            for name in files:
                path = self.transcripts_dir / name
                for record in _stream_jsonl_file(path):
                    if record.get("type") == record_type:
                        records.append(record)
            records.sort(key=lambda r: r.get("seq", 0))
            self._type_index[record_type] = records
            self._loaded_types.add(record_type)

    def _records_of_type(self, record_type: str) -> List[Dict[str, Any]]:
        self._load_type(record_type)
        return self._type_index.get(record_type, [])

    def expect(self, record_type: str, task_id: Optional[int] = None) -> Dict[str, Any]:
        with self._lock:
            type_records = self._records_of_type(record_type)
            pos = self.positions.get(record_type, 0)
            if pos >= len(type_records):
                raise ReplayDivergence("MISSING_DEP", {"type": record_type, "pos": pos})
            record = type_records[pos]
            self.positions[record_type] = pos + 1
            return record

    def has_more(self, record_type: str) -> bool:
        with self._lock:
            pos = self.positions.get(record_type, 0)
            return pos < len(self._records_of_type(record_type))


def load_transcripts(
    transcripts_dir: Path,
    use_streaming: Optional[bool] = None,
) -> TranscriptIndex:
    if use_streaming is None:
        size = _estimate_transcript_size(transcripts_dir)
        use_streaming = size >= STREAMING_THRESHOLD_BYTES

    if use_streaming:
        return StreamingTranscriptIndex(transcripts_dir)  # type: ignore[return-value]

    records: List[Dict[str, Any]] = []
    for name in [
        "events.jsonl",
        "deps_http.jsonl",
        "deps_sql.jsonl",
        "deps_redis.jsonl",
        "deps_grpc.jsonl",
        "deps_aws.jsonl",
        "deps_mongodb.jsonl",
        "time_rng.jsonl",
    ]:
        path = transcripts_dir / name
        records.extend(_read_jsonl_file(path))
    records.sort(key=lambda r: r.get("seq", 0))

    # Load async task graph if present (v2.0 bundles)
    task_graph = _read_jsonl_file(transcripts_dir / "async_task_graph.jsonl")
    task_graph.sort(key=lambda r: r.get("seq", 0))

    return TranscriptIndex(records, task_graph=task_graph)
