"""Index and read Medulla DObject recordings.

Supports:

- standard layout: ``dobject/**/*.log`` + ``dobject_recording/**/*.dorec``
- recovered layout: ``dobject/all/indices.log`` + ``dobject_recording/data.bin``
- either as an extracted directory or a zip containing those paths
"""

from __future__ import annotations

import re
import struct
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO, Iterator


RECORD_RE = re.compile(
    r"^(?:\[(?P<log_time>[^]]+)\])?>?\s*DObject `(?P<name>[^`]+)` post "
    r"len=(?P<len>\d+)B, id:(?P<id>[0-9A-Fa-f]+), tic:(?P<tic>\d+), "
    r"@(?P<file>[^:]+):(?P<offset>\d+)"
)


@dataclass(frozen=True)
class RecordRef:
    sequence: int
    object_name: str
    log_time: str
    source_log: str
    source_dorec: str
    source_offset: int
    payload_length: int
    log_record_id: str
    dotnet_ticks: int


class _ZipStoredMemberIO:
    """Random-access reader for a ZIP_STORED member via the underlying zip file.

    ``ZipExtFile.seek`` on multi-GB members is far too slow for per-record reads.
    """

    def __init__(self, zip_path: Path, member_name: str, data_offset: int, data_size: int):
        self._path = zip_path
        self._member_name = member_name
        self._data_offset = data_offset
        self._data_size = data_size
        self._fh = zip_path.open("rb")
        self._pos = 0

    def seek(self, offset: int, whence: int = 0) -> int:
        if whence == 0:
            self._pos = offset
        elif whence == 1:
            self._pos += offset
        elif whence == 2:
            self._pos = self._data_size + offset
        else:
            raise ValueError(f"invalid whence: {whence}")
        if self._pos < 0:
            raise ValueError("negative seek")
        return self._pos

    def read(self, size: int = -1) -> bytes:
        if size is None or size < 0:
            size = self._data_size - self._pos
        if size <= 0 or self._pos >= self._data_size:
            return b""
        size = min(size, self._data_size - self._pos)
        self._fh.seek(self._data_offset + self._pos)
        data = self._fh.read(size)
        self._pos += len(data)
        return data

    def close(self) -> None:
        self._fh.close()


def _zip_stored_member_offset(zip_path: Path, info: zipfile.ZipInfo) -> int:
    if info.compress_type != zipfile.ZIP_STORED:
        raise RuntimeError(
            f"member {info.filename!r} is compressed (type={info.compress_type}); "
            "extract it first or store uncompressed"
        )
    with zip_path.open("rb") as handle:
        handle.seek(info.header_offset)
        header = handle.read(30)
        if len(header) != 30 or header[:4] != b"PK\x03\x04":
            raise RuntimeError(f"bad local zip header for {info.filename!r}")
        name_len, extra_len = struct.unpack("<HH", header[26:30])
        return info.header_offset + 30 + name_len + extra_len


@dataclass
class DlogSource:
    """Opened dlog directory or recovered zip."""

    label: str
    directory: Path | None = None
    zip_path: Path | None = None
    _zip: zipfile.ZipFile | None = None
    _log_cache: dict[str, str] | None = None
    _member_offsets: dict[str, tuple[int, int]] | None = None

    def close(self) -> None:
        if self._zip is not None:
            self._zip.close()
            self._zip = None

    def __enter__(self) -> "DlogSource":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def iter_log_texts(self) -> Iterator[tuple[str, str]]:
        if self.zip_path is not None:
            assert self._zip is not None
            if self._log_cache is None:
                self._log_cache = {}
                names = sorted(
                    name
                    for name in self._zip.namelist()
                    if name.replace("\\", "/").startswith("dobject/")
                    and name.replace("\\", "/").endswith(".log")
                )
                for name in names:
                    key = name.replace("\\", "/")
                    self._log_cache[key] = self._zip.read(name).decode("utf-8", errors="replace")
            for name, text in self._log_cache.items():
                yield name, text
            return
        assert self.directory is not None
        for log_path in sorted((self.directory / "dobject").rglob("*.log")):
            relative = log_path.relative_to(self.directory).as_posix()
            yield relative, log_path.read_text(encoding="utf-8", errors="replace")

    def open_recording(self, name: str) -> tuple[object, BinaryIO]:
        """Return (owner, binary stream) supporting seek/read of one recording member."""

        base = Path(name).name
        if self.zip_path is not None:
            assert self._zip is not None
            candidates = [
                n
                for n in self._zip.namelist()
                if Path(n.replace("\\", "/")).name.casefold() == base.casefold()
                and "dobject_recording/" in n.replace("\\", "/")
            ]
            if not candidates:
                alt = name.replace("\\", "/")
                if alt in self._zip.namelist():
                    candidates = [alt]
                elif f"dobject_recording/{base}" in self._zip.namelist():
                    candidates = [f"dobject_recording/{base}"]
            if not candidates:
                raise FileNotFoundError(f"missing recording in zip: {name}")
            if len(candidates) > 1:
                raise RuntimeError(f"ambiguous recording in zip {name}: {candidates}")
            member = candidates[0].replace("\\", "/")
            if self._member_offsets is None:
                self._member_offsets = {}
            if member not in self._member_offsets:
                info = self._zip.getinfo(member)
                self._member_offsets[member] = (
                    _zip_stored_member_offset(self.zip_path, info),
                    info.file_size,
                )
            data_offset, data_size = self._member_offsets[member]
            stream = _ZipStoredMemberIO(self.zip_path, member, data_offset, data_size)
            return stream, stream

        assert self.directory is not None
        index = index_dorec_files(self.directory)
        if base.casefold() == "data.bin":
            path = self.directory / "dobject_recording" / "data.bin"
            if not path.is_file():
                matches = list((self.directory / "dobject_recording").rglob("data.bin"))
                if not matches:
                    raise FileNotFoundError(f"missing recording file: {name}")
                path = matches[0]
            stream = path.open("rb")
            return stream, stream
        path = choose_dorec(index, name)
        stream = path.open("rb")
        return stream, stream


def open_dlog_source(value: Path | str) -> DlogSource:
    path = Path(value).expanduser().resolve()
    if path.is_file() and path.suffix.lower() == ".zip":
        zf = zipfile.ZipFile(path, "r")
        names = {n.replace("\\", "/") for n in zf.namelist()}
        has_log = any(n.startswith("dobject/") and n.endswith(".log") for n in names)
        has_rec = any(n.startswith("dobject_recording/") for n in names)
        if not (has_log and has_rec):
            zf.close()
            raise FileNotFoundError(f"{path} is not a recovered/standard dlog zip")
        return DlogSource(label=str(path), zip_path=path, _zip=zf)

    root = path
    if not ((root / "dobject").is_dir() and (root / "dobject_recording").is_dir()):
        child = root / "dlog"
        if (child / "dobject").is_dir() and (child / "dobject_recording").is_dir():
            root = child
        else:
            raise FileNotFoundError(f"{path} does not contain dobject and dobject_recording")
    return DlogSource(label=str(root), directory=root)


def resolve_dlog_root(value: Path | str) -> Path:
    """Backward-compatible helper: directory roots only (not zip)."""

    source = open_dlog_source(value)
    try:
        if source.directory is None:
            raise FileNotFoundError(
                f"{value} is a zip; use open_dlog_source()/iter_payloads_from_source()"
            )
        return source.directory
    finally:
        source.close()


def discover_records_from_source(
    source: DlogSource,
    object_name: str,
    *,
    host_ticks_min: int | None = None,
    host_ticks_max: int | None = None,
) -> list[RecordRef]:
    pending: list[tuple[str, str, str, int, int, str, int, str]] = []
    name_key = object_name.casefold()
    for relative_log, text in source.iter_log_texts():
        for line in text.splitlines():
            match = RECORD_RE.search(line.strip())
            if not match or match.group("name").casefold() != name_key:
                continue
            ticks = int(match.group("tic"))
            if host_ticks_min is not None and ticks < host_ticks_min:
                continue
            if host_ticks_max is not None and ticks > host_ticks_max:
                continue
            pending.append(
                (
                    match.group("name"),
                    match.group("log_time") or "",
                    relative_log,
                    int(match.group("offset")),
                    int(match.group("len")),
                    match.group("id").upper(),
                    ticks,
                    match.group("file"),
                )
            )
    pending.sort(key=lambda item: (item[6], item[7].casefold(), item[3]))
    seen: set[tuple[str, int, int]] = set()
    records: list[RecordRef] = []
    for item in pending:
        key = (item[7].casefold(), item[3], item[6])
        if key in seen:
            continue
        seen.add(key)
        records.append(
            RecordRef(
                sequence=len(records),
                object_name=item[0],
                log_time=item[1],
                source_log=item[2],
                source_dorec=item[7],
                source_offset=item[3],
                payload_length=item[4],
                log_record_id=item[5],
                dotnet_ticks=item[6],
            )
        )
    return records


def discover_records(dlog_root: Path, object_name: str) -> list[RecordRef]:
    with open_dlog_source(dlog_root) as source:
        return discover_records_from_source(source, object_name)


def index_dorec_files(dlog_root: Path) -> dict[str, list[Path]]:
    result: dict[str, list[Path]] = {}
    recording = dlog_root / "dobject_recording"
    if not recording.is_dir():
        return result
    for path in recording.rglob("*"):
        if path.is_file() and path.suffix.lower() in {".dorec", ".bin"}:
            result.setdefault(path.name.casefold(), []).append(path)
    return result


def choose_dorec(index: dict[str, list[Path]], name: str) -> Path:
    matches = index.get(Path(name).name.casefold(), [])
    if not matches:
        raise FileNotFoundError(f"missing recording file: {name}")
    if len(matches) > 1:
        raise RuntimeError(f"ambiguous recording file {name}: {matches}")
    return matches[0]


def read_exact(stream: BinaryIO, size: int) -> bytes:
    data = stream.read(size)
    if len(data) != size:
        raise EOFError(f"expected {size} bytes, got {len(data)}")
    return data


def _read_payload_at(stream: BinaryIO, record: RecordRef) -> bytes:
    stream.seek(record.source_offset)
    name_length = read_exact(stream, 1)[0]
    name = read_exact(stream, name_length).decode("ascii")
    ticks = struct.unpack("<q", read_exact(stream, 8))[0]
    id_length = read_exact(stream, 1)[0]
    id_bytes = read_exact(stream, id_length)
    payload_length = struct.unpack("<i", read_exact(stream, 4))[0]
    payload = read_exact(stream, payload_length)
    try:
        record_id = id_bytes.decode("ascii")
    except UnicodeDecodeError:
        record_id = id_bytes.hex().upper()
    if name != record.object_name:
        raise ValueError(f"name mismatch: log={record.object_name}, dorec={name}")
    if ticks != record.dotnet_ticks:
        raise ValueError(f"tick mismatch: log={record.dotnet_ticks}, dorec={ticks}")
    if payload_length != record.payload_length:
        raise ValueError(f"payload mismatch: log={record.payload_length}, dorec={payload_length}")
    if record_id.upper() != record.log_record_id.upper():
        raise ValueError(f"record id mismatch: log={record.log_record_id}, dorec={record_id}")
    return payload


def iter_payloads_from_source(
    source: DlogSource,
    object_name: str,
    *,
    host_ticks_min: int | None = None,
    host_ticks_max: int | None = None,
) -> Iterator[tuple[RecordRef, bytes]]:
    records = discover_records_from_source(
        source,
        object_name,
        host_ticks_min=host_ticks_min,
        host_ticks_max=host_ticks_max,
    )
    if not records:
        return
    open_files: dict[str, BinaryIO] = {}
    try:
        for record in records:
            key = Path(record.source_dorec).name.casefold()
            stream = open_files.get(key)
            if stream is None:
                _owner, stream = source.open_recording(record.source_dorec)
                open_files[key] = stream
            yield record, _read_payload_at(stream, record)
    finally:
        for stream in open_files.values():
            stream.close()


def iter_payloads(
    dlog_root: Path | str,
    object_name: str,
    *,
    host_ticks_min: int | None = None,
    host_ticks_max: int | None = None,
) -> Iterator[tuple[RecordRef, bytes]]:
    with open_dlog_source(dlog_root) as source:
        yield from iter_payloads_from_source(
            source,
            object_name,
            host_ticks_min=host_ticks_min,
            host_ticks_max=host_ticks_max,
        )