# Standard library imports
import logging
import threading
from collections import deque
from datetime import date, datetime, timedelta
from queue import Empty, Queue
from time import monotonic
from typing import TYPE_CHECKING, Any, TextIO

# Third party imports
from orjson import JSONDecodeError, dumps, loads
from pydantic import ConfigDict
from pydantic.dataclasses import dataclass

# First party imports
from aeth_ext.central_log_server.protocol import payload_to_record, record_to_payload
from aeth_ext.errors import handle_fatal_exc_sync
from aeth_ext.logging.bases import TaggedLogRecord
from aeth_ext.settings import BaseSettings
from aeth_ext.types import IsPydanticSlots

if TYPE_CHECKING:
  # Standard library imports
  from collections.abc import Iterator
  from pathlib import Path

logger = logging.getLogger(__name__)

__all__ = ["EmergencyHistoryWriter", "HistoryEntry", "RecordHistoryBuffer"]

settings = BaseSettings.get_settings()

# How long EmergencyHistoryWriter's run loop waits for the next entry before
# treating the queue as idle and closing the open file handle.
_IDLE_CLOSE_TIMEOUT = 2.0

# How far back from the end of a history file `_repair_truncated_tail` will look
# for the last newline. History files are unbounded, but only the final line can
# ever be in question, so reading the whole file to inspect its last byte would
# be wasteful. Generous enough that a realistic record always fits.
_MAX_TAIL_SCAN_BYTES = 1024 * 1024

pyd_config = ConfigDict(arbitrary_types_allowed=True)


@dataclass(config=pyd_config, slots=True)
class HistoryEntry(IsPydanticSlots):
  """A single emitted record, kept around so it can be replayed on reconnect.

  ``persisted`` tracks whether this entry has already been written to a
  history file (e.g. by :class:`EmergencyHistoryWriter`) so that
  :class:`RecordHistoryBuffer`'s own lazy spill does not write it a second
  time once it is evicted from memory.
  """

  id: int
  created: float
  record: TaggedLogRecord
  persisted: bool = False


def _history_file_for_date(history_dir: Path, day: date) -> Path:
  return history_dir / f"{day:%Y-%m-%d}.jsonl"


def _format_entry_line(entry: HistoryEntry) -> str:
  """Serialise *entry* to the JSON line written to a history file."""
  return dumps(
    {
      "id": entry.id,
      "created": entry.created,
      "record": record_to_payload(entry.record),
    },
    default=str,
  ).decode()


def iter_entries(path: Path) -> Iterator[HistoryEntry]:
  """Yield every :class:`HistoryEntry` in the JSONL file at *path*, in order.

  Malformed lines (e.g. a truncated final line left by a crash mid-write) are
  logged and skipped rather than aborting the whole read.
  """
  if not path.exists():
    return
  with path.open("r", encoding="utf-8") as fh:
    for line in fh:
      stripped_line = line.strip()
      if not stripped_line:
        continue
      try:
        data: dict[str, object] = loads(stripped_line)
        record = payload_to_record(data["record"])  # pyright: ignore[reportArgumentType]
      except JSONDecodeError, KeyError, ValueError, TypeError:
        logger.warning("Skipping malformed history line in %s", path, exc_info=True)
        continue
      yield HistoryEntry(id=int(data["id"]), created=float(data["created"]), record=record, persisted=True)  # pyright: ignore[reportArgumentType]


def _approx_entry_size(record: TaggedLogRecord) -> int:
  """A cheap size estimate used only to bound memory, not an exact size.

  Avoids pickling every record just to measure it; a fixed per-record
  overhead plus the message length is close enough for a memory cap.
  """
  try:
    message_len = len(record.getMessage())
  except Exception:  # noqa: BLE001
    message_len = 0
  return message_len + 512


class RecordHistoryBuffer:
  """An in-memory, append-only record of everything emitted (sent or not).

  Bounded by three independent thresholds - record count, approximate byte
  size, and time since the last flush - so a crash can lose at most the last
  ``max_age`` seconds of records rather than an unbounded amount. Whichever
  threshold trips first flushes *every* currently-buffered entry to its
  date-segregated JSONL history file and clears the in-memory buffer, so
  :meth:`find_after` only ever needs to search memory for the most recent
  contiguous tail and disk for anything older.

  Each instance gets its own subdirectory under :attr:`history_root`, keyed
  by *program_name*: the drainer classes are primarily used to submit records
  generated by a separate program (a subprocess, a different package) to the
  log server, so that program's history must stay in its own files rather
  than interleaving with whatever the current process's own logging is
  already writing under the same date-keyed name.
  """

  history_root = settings.persisted_dir_loc / "client_log_history"

  def __init__(
    self,
    program_name: str,
    max_records: int = 50_000,
    max_bytes: int = 64 * 1024 * 1024,
    max_age: float = 300.0,
  ) -> None:
    self.history_dir = self.history_root / program_name
    self.history_dir.mkdir(parents=True, exist_ok=True)
    self._max_records = max_records
    self._max_bytes = max_bytes
    self._max_age = max_age
    self._entries: deque[HistoryEntry] = deque()
    self._approx_bytes = 0
    self._last_flush_monotonic = monotonic()
    # Guards `_entries` and the whole write path. All three client classes call
    # `append` from arbitrary threads, so `_flush_to_disk`'s drain loop could
    # already interleave today; write-through mode (below) would turn that from
    # rare into per-record, and two threads appending to the same file at once
    # interleave partial JSON lines *mid-file* -- damage `_repair_truncated_tail`
    # cannot fix, since it only ever inspects the trailing segment.
    # Reentrant because `append` holds it across `_maybe_flush` -> `_flush_to_disk`.
    self._lock = threading.RLock()
    # Flipped by `begin_shutdown`. Never read without also consulting the lock
    # *except* in `begin_shutdown` itself -- see there for why.
    self._shutting_down = False
    # Flipped once, by whichever of `append`/`flush` is first to observe
    # `_shutting_down` after it goes true. Marks that the one-time catch-up
    # drain (see `append`) has happened and `_entries`/`_flush_to_disk` are
    # retired for the rest of this buffer's life -- every later write-through
    # entry goes straight to disk via `_write_entry_through` instead.
    self._shutdown_flushed = False
    # The file handle `_write_entry_through` is currently writing through to,
    # held open across calls for as long as it keeps targeting the same date
    # file -- graceful shutdown is exactly when execution time is scarcest, so
    # paying one `open` per write-through entry (as opposed to one per burst)
    # would be spending that budget in the worst possible place. Closed and
    # cleared by `close`.
    self._write_through_fh: TextIO | None = None
    self._write_through_path: Path | None = None
    self._repair_truncated_tail()

  def _repair_truncated_tail(self) -> None:
    """Repair a final record left half-written by a previous run (D-I7).

    Only today's file is checked -- it is the only one a fresh instance is about
    to append to; other dates are read-only lookup targets for
    :meth:`find_after`.

    Every successful write appends a complete JSON line *plus* a trailing
    newline, so a file not ending in ``\\n`` is by construction evidence that
    something was cut short. The trailing segment is then parsed to decide which
    of two very different repairs applies -- deleting a record that turns out to
    be intact would itself be data loss.

    Reads a bounded tail rather than the whole file: these can grow without
    limit, and only the last line is ever in question.

    Best-effort throughout. Any :exc:`OSError` is logged and swallowed, since
    construction succeeding matters more than the repair.
    """
    path = _history_file_for_date(self.history_dir, datetime.now(tz=settings.tz).date())
    try:
      size = path.stat().st_size if path.exists() else 0
      if size == 0:
        return

      read_size = min(size, _MAX_TAIL_SCAN_BYTES)
      with path.open("rb") as fh:
        fh.seek(size - read_size)
        tail = fh.read(read_size)

      if tail.endswith(b"\n"):
        return  # well-formed; nothing to do

      newline_at = tail.rfind(b"\n")
      if newline_at == -1 and read_size < size:
        # The final line is longer than the scan window, so its true start is
        # somewhere we did not read. Truncating would need that boundary, and
        # guessing it risks discarding good records -- leave the file alone.
        logger.warning("History file %s has an unterminated final line longer than %d bytes; leaving as-is", path, read_size)
        return

      suspect = tail[newline_at + 1 :]
      try:
        loads(suspect)
      except JSONDecodeError:
        # A genuine partial write. Report the length rather than the content:
        # a truncated write need not even be valid UTF-8.
        truncate_at = size - read_size + newline_at + 1
        with path.open("r+b") as fh:
          fh.truncate(truncate_at)
        logger.warning("Discarded a %d-byte partial record at the end of %s", len(suspect), path)
      else:
        # A complete record that only lost its newline. Should not happen given
        # the write pattern, but appending one byte is far cheaper than assuming
        # it cannot and dropping a valid record.
        with path.open("ab") as fh:
          fh.write(b"\n")
        logger.warning("Restored the missing trailing newline on %s", path)
    except OSError:
      logger.exception("Could not repair the trailing record in %s", path)

  def begin_shutdown(self) -> None:
    """Switch to write-through mode; safe to call from a signal handler (D-I6).

    Registered for :attr:`~aeth_ext.errors.ShutdownPhase.INTERRUPT`, so this runs
    on the main thread between bytecodes and must obey that phase's rules: it
    does exactly one atomic attribute store, and **nothing else**.

    In particular it deliberately does *not* flush, and does not take
    :attr:`_lock`. Flushing here would re-enter :meth:`_flush_to_disk`'s drain
    loop if the interrupted thread was already inside it, losing or duplicating
    entries; taking the lock would deadlock outright for the same reason. The
    backlog is not stranded by that choice -- whichever of :meth:`append` or
    :meth:`flush` next observes :attr:`_shutting_down` performs the one-time
    catch-up drain itself (see :meth:`append`).
    """
    self._shutting_down = True

  def flush(self) -> None:
    """Force every buffered entry to disk, ignoring the usual thresholds.

    A no-op once the shutdown catch-up drain has already run (see
    :meth:`append`) -- at that point :attr:`_entries` is retired for good and
    there is nothing left in it to flush.
    """
    with self._lock:
      if self._shutting_down and self._shutdown_flushed:
        return
      self._flush_to_disk()
      self._last_flush_monotonic = monotonic()
      if self._shutting_down:
        self._shutdown_flushed = True

  def append(self, entry: HistoryEntry) -> None:
    with self._lock:
      if self._shutting_down:
        if self._shutdown_flushed:
          # Fast path: the one-time catch-up drain below already ran, so
          # `_entries` is retired -- write straight to disk without spinning
          # up `_flush_to_disk`'s multi-entry drain loop for what is, from
          # here on, always exactly one entry.
          self._write_entry_through(entry)
        else:
          # First write-through append: catch up whatever `begin_shutdown`
          # left buffered, plus this entry, in one drain, then retire
          # `_entries` for the rest of this buffer's life.
          self._entries.append(entry)
          self._approx_bytes += _approx_entry_size(entry.record)
          self._flush_to_disk()
          self._last_flush_monotonic = monotonic()
          self._shutdown_flushed = True
      else:
        self._entries.append(entry)
        self._approx_bytes += _approx_entry_size(entry.record)
        self._maybe_flush()

  def _write_entry_through(self, entry: HistoryEntry) -> None:
    """Persist a single entry directly, bypassing :attr:`_entries` entirely.

    Caller must hold :attr:`_lock`. Only ever reached once the shutdown
    catch-up drain has run (:attr:`_shutdown_flushed` is set), so there is
    never more than one entry to write here.

    :attr:`_write_through_fh` is reused across calls rather than reopened per
    entry -- write-through mode exists for graceful shutdown, where CPU time
    is the scarcest resource, so paying an ``open`` per record would be
    spending that budget on exactly the wrong thing. Each write is still
    flushed individually so no record is stranded in a buffer if the process
    is killed before :meth:`close` runs.
    """
    day = datetime.fromtimestamp(entry.created, tz=settings.tz).date()
    path = _history_file_for_date(self.history_dir, day)
    if path != self._write_through_path:
      self._close_write_through()
      try:
        self._write_through_fh = path.open("a", encoding="utf-8")
      except OSError:
        logger.exception("Failed to open history file %s for entry %s", path, entry.id)
        return
      self._write_through_path = path
    try:
      assert self._write_through_fh is not None
      self._write_through_fh.write(_format_entry_line(entry) + "\n")
      self._write_through_fh.flush()
    except OSError:
      logger.exception("Failed to spill history entry %s to disk", entry.id)

  def _close_write_through(self) -> None:
    """Close and clear :attr:`_write_through_fh`, if open. Caller must hold :attr:`_lock`."""
    if self._write_through_fh is not None:
      try:
        self._write_through_fh.close()
      except OSError:
        pass
      self._write_through_fh = None
      self._write_through_path = None

  def close(self) -> None:
    """Release the write-through file handle held open by :meth:`_write_entry_through`.

    A no-op if write-through mode was never entered. Meant to run at final
    teardown, after the last possible :meth:`append`/:meth:`flush` call --
    there is no way back into non-write-through mode, so nothing needs this
    handle again once it is closed.
    """
    with self._lock:
      self._close_write_through()

  def _maybe_flush(self) -> None:
    """Flush if any threshold has tripped. Caller must hold :attr:`_lock`."""
    now = monotonic()
    if (
      len(self._entries) < self._max_records
      and self._approx_bytes < self._max_bytes
      and (now - self._last_flush_monotonic) < self._max_age
    ):
      return
    self._flush_to_disk()
    self._last_flush_monotonic = now

  def _flush_to_disk(self) -> None:
    """Drain the buffer to date-segregated files. Caller must hold :attr:`_lock`."""
    open_files: dict[Path, Any] = {}
    try:
      while self._entries:
        entry = self._entries.popleft()
        if entry.persisted:
          continue
        day = datetime.fromtimestamp(entry.created, tz=settings.tz).date()
        path = _history_file_for_date(self.history_dir, day)
        if path not in open_files:
          try:
            open_files[path] = path.open("a", encoding="utf-8")
          except OSError:
            logger.exception("Failed to open history file %s for entry %s", path, entry.id)
            continue
        try:
          open_files[path].write(_format_entry_line(entry) + "\n")
        except Exception:
          logger.exception("Failed to spill history entry %s to disk", entry.id)
    finally:
      for fh in open_files.values():
        try:
          fh.close()
        except OSError:
          pass
    self._approx_bytes = 0

  def find_after(self, last_id: int | None, hint_created: float | None) -> tuple[HistoryEntry, ...] | None:
    """Return every entry with ``id > last_id``, or ``None`` if unrecoverable.

    ``last_id`` of ``None`` means the server has never seen this program
    before, so everything currently retained (memory and, in principle, disk)
    is replayed. ``None`` is returned only when ``last_id`` is a real id that
    cannot be located anywhere (memory or the probed history files) - a
    genuine, unrecoverable gap rather than "nothing to resend".
    """
    if last_id is None:
      return tuple(self._entries)

    if self._entries and self._entries[0].id <= last_id:
      return tuple(e for e in self._entries if e.id > last_id)

    if self._entries and last_id > self._entries[-1].id:
      # Already caught up (or the ack is stale); nothing to resend.
      return ()

    disk_entries = self._search_disk(last_id, hint_created)
    if disk_entries is None:
      return None
    return (*disk_entries, *self._entries)

  def _search_disk(self, last_id: int, hint_created: float | None) -> tuple[HistoryEntry, ...] | None:
    """Scan history files near *hint_created* for everything after *last_id*.

    Probes the hinted date plus the adjacent day on either side to absorb
    midnight-rollover ambiguity between when a record was created and which
    date file it landed in. This is a bounded, practical window intended for
    the short reconnect gaps this mechanism targets, not an exhaustive scan of
    all retained history.
    """
    hint_date = (
      datetime.fromtimestamp(hint_created, tz=settings.tz).date() if hint_created is not None else datetime.now(tz=settings.tz).date()
    )
    candidate_dates = (hint_date - timedelta(days=1), hint_date, hint_date + timedelta(days=1))

    found = False
    collected: list[HistoryEntry] = []
    for day in candidate_dates:
      for entry in iter_entries(_history_file_for_date(self.history_dir, day)):
        if found:
          collected.append(entry)
        elif entry.id == last_id:
          found = True
        elif entry.id > last_id:
          # last_id itself wasn't found verbatim (e.g. it belongs to a file
          # outside the probed window); treat the first newer entry seen as
          # the resume point so nothing after it is missed.
          found = True
          collected.append(entry)

    return tuple(collected) if found else None


class EmergencyHistoryWriter:
  """Eagerly persists every submitted record while the server is unreachable.

  Started only once the caller's send-failure heuristic trips, and stopped as
  soon as the connection recovers, so the extra standing thread is paid for
  only while genuinely needed. Runs on its own daemon thread with a simple
  FIFO queue so submitting a record from :meth:`~logging.Handler.emit` never
  blocks on disk IO.

  The file handle for the active date's history file is held open across
  consecutive writes and only closed once the queue has been idle for
  ``_IDLE_CLOSE_TIMEOUT`` seconds or the date rolls over, so a burst of
  records pays for one ``open`` rather than one per record.
  """

  def __init__(self, history_dir: Path) -> None:
    self._history_dir = history_dir
    self._history_dir.mkdir(parents=True, exist_ok=True)
    self._queue: Queue[HistoryEntry | None] = Queue()
    self._thread = threading.Thread(target=self._run, name="log-emergency-writer", daemon=True)
    self._thread.start()

  def submit(self, entry: HistoryEntry) -> None:
    self._queue.put(entry)

  def _close_handle(self, fh: Any) -> None:
    """Close *fh*, silently swallowing any ``OSError``."""
    if fh is not None:
      try:
        fh.close()
      except OSError:
        pass

  def _switch_file(self, current_fh: Any, new_path: Path) -> Any:
    """Close *current_fh* (if open) and open *new_path* for append.

    Returns the new file handle, or ``None`` if the open fails.
    """
    self._close_handle(current_fh)
    try:
      return new_path.open("a", encoding="utf-8")
    except OSError:
      logger.exception("Emergency history writer failed to open %s", new_path)
      return None

  @handle_fatal_exc_sync
  def _run(self) -> None:
    current_fh = None
    current_path = None

    while True:
      try:
        entry = self._queue.get(timeout=_IDLE_CLOSE_TIMEOUT)
      except Empty:
        # Queue went idle: release the handle so it is not held open
        # indefinitely between bursts, then block until the next entry.
        self._close_handle(current_fh)
        current_fh = None
        current_path = None
        entry = self._queue.get()

      if entry is None:
        break

      day = datetime.fromtimestamp(entry.created, tz=settings.tz).date()
      path = _history_file_for_date(self._history_dir, day)

      if path != current_path:
        current_fh = self._switch_file(current_fh, path)
        current_path = path if current_fh is not None else None

      if current_fh is None:
        continue

      try:
        current_fh.write(_format_entry_line(entry) + "\n")
        current_fh.flush()
      except OSError:
        logger.exception("Emergency history writer failed to persist record %s", entry.id)

    self._close_handle(current_fh)

  def close(self) -> None:
    self._queue.put(None)
    self._thread.join(timeout=5.0)
