"""
JSONL persistence layer for the database.
"""

import copy
import logging
import os
import signal
from collections import deque
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from uuid import UUID

import aiofiles
import jsondiff
import msgspec

from paskia.db.logging import log_change
from paskia.db.migrations import DBVER, apply_all_migrations
from paskia.db.structs import DB, SessionContext

_logger = logging.getLogger(__name__)

# Default database path
DB_PATH_DEFAULT = "paskia.jsonl"


class _ChangeRecord(msgspec.Struct, omit_defaults=True):
    """A single change record in the JSONL file."""

    ts: datetime
    a: str  # action - describes the operation (e.g., "migrate", "login", "create_user")
    v: int  # schema version after this change
    u: str | None = None  # user UUID who performed the action (None for system)
    diff: dict = {}


# msgspec encoder for change records
_change_encoder = msgspec.json.Encoder()


def compute_diff(previous: dict, current: dict) -> dict | None:
    """Compute JSON diff between two states.

    Args:
        previous: Previous state (JSON-compatible dict)
        current: Current state (JSON-compatible dict)

    Returns:
        The diff, or None if no changes
    """
    diff = jsondiff.diff(previous, current, marshal=True)
    return diff if diff else None


def create_change_record(
    action: str, version: int, diff: dict, user: str | None = None
) -> _ChangeRecord:
    """Create a change record for persistence."""
    return _ChangeRecord(
        ts=datetime.now(UTC),
        a=action,
        v=version,
        u=user,
        diff=diff,
    )


# Actions that are allowed to create a new database file
_BOOTSTRAP_ACTIONS = frozenset({"bootstrap"})

# Flag to prevent duplicate error messages on fatal flush failure
_flush_failed = False


async def flush_changes(
    db_path: Path,
    pending_changes: deque[_ChangeRecord],
) -> None:
    """Write all pending changes to disk.

    Args:
        db_path: Path to the JSONL database file
        pending_changes: Queue of pending change records (will be cleared on success)

    On failure, logs an error and sends SIGTERM to trigger graceful shutdown.
    """
    global _flush_failed
    if _flush_failed or not pending_changes:
        return

    if not db_path.exists():
        first_action = pending_changes[0].a
        if first_action not in _BOOTSTRAP_ACTIONS:
            _logger.error(
                "Refusing to create database file with action '%s' - "
                "only bootstrap can create a new database",
                first_action,
            )
            _flush_failed = True
            os.kill(os.getpid(), signal.SIGTERM)
            return

    changes_to_write = list(pending_changes)

    try:
        lines = [_change_encoder.encode(change) for change in changes_to_write]
        if not lines:
            pending_changes.clear()
            return

        async with aiofiles.open(db_path, "ab") as f:
            await f.write(b"\n".join(lines) + b"\n")
        pending_changes.clear()
    except OSError as e:
        _logger.error("Failed to flush database: %s", e)
        _flush_failed = True
        os.kill(os.getpid(), signal.SIGTERM)


class JsonlStore:
    """JSONL persistence layer for a DB instance."""

    def __init__(self, db: DB, db_path: str = DB_PATH_DEFAULT):
        self.db: DB = db
        self.db_path = Path(db_path)
        self._previous_builtins: dict[str, Any] = {}
        self._pending_changes: deque[_ChangeRecord] = deque()
        self._current_action: str = "system"
        self._current_user: str | None = None
        self._in_transaction: bool = False
        self._transaction_snapshot: dict[str, Any] | None = None
        self._current_version: int = DBVER  # Schema version for new databases

    async def load(
        self, db_path: str | None = None, *, rp_id: str = "localhost"
    ) -> None:
        """Load data from JSONL change log."""
        if db_path is not None:
            self.db_path = Path(db_path)
        self._rp_id = rp_id
        if not self.db_path.exists():
            return

        # Replay change log to reconstruct state
        data_dict: dict = {}
        try:
            async with aiofiles.open(self.db_path, "rb") as f:
                content = await f.read()
            for line_num, line in enumerate(content.split(b"\n"), 1):
                line = line.strip()
                if not line:
                    continue
                try:
                    change = msgspec.json.decode(line)
                    data_dict = jsondiff.patch(data_dict, change["diff"], marshal=True)
                    self._current_version = change.get("v", 0)
                except Exception as e:
                    raise ValueError(f"Error parsing line {line_num}: {e}")
        except OSError as e:
            raise SystemExit(f"Failed to load database: {e}")
        except (ValueError, msgspec.DecodeError) as e:
            raise SystemExit(f"Failed to load database: {e}")

        if not data_dict:
            return

        # Set previous state for diffing (will be updated by _queue_change)
        self._previous_builtins = copy.deepcopy(data_dict)

        # Callback to persist each migration
        async def persist_migration(
            action: str, new_version: int, current: dict
        ) -> None:
            self._current_version = new_version
            self._queue_change(action, new_version, current)

        # Apply schema migrations one at a time
        await apply_all_migrations(
            data_dict, self._current_version, persist_migration, rp_id=rp_id
        )

        # Decode to msgspec struct
        decoder = msgspec.json.Decoder(DB)
        self.db = decoder.decode(msgspec.json.encode(data_dict))
        self.db._store = self

        # Normalize via msgspec round-trip (handles omit_defaults etc.)
        # This ensures _previous_builtins matches what msgspec would produce
        normalized_dict = msgspec.to_builtins(self.db)
        await persist_migration(
            "migrate:msgspec", self._current_version, normalized_dict
        )

    def _queue_change(
        self, action: str, version: int, current: dict, user: str | None = None
    ) -> None:
        """Queue a change record and log it.

        Args:
            action: The action name for the change record
            version: The schema version for the change record
            current: The current state as a plain dict
            user: Optional user UUID who performed the action
        """
        diff = compute_diff(self._previous_builtins, current)
        if not diff:
            return
        self._pending_changes.append(create_change_record(action, version, diff, user))

        # Log the change with user display name if available
        user_display = None
        if user:
            try:
                user_uuid = UUID(user)
                if user_uuid in self.db.users:
                    user_display = self.db.users[user_uuid].display_name
            except (ValueError, KeyError):
                user_display = user

        log_change(action, diff, user_display, self._previous_builtins, self.db)
        self._previous_builtins = copy.deepcopy(current)

    @contextmanager
    def transaction(
        self,
        action: str,
        ctx: SessionContext | None = None,
        *,
        user: str | None = None,
    ):
        """Wrap writes in transaction. Queues change on successful exit.

        Args:
            action: Describes the operation (e.g., "Created user", "Login")
            ctx: Session context of user performing the action (None for system operations)
            user: User UUID string (alternative to ctx when full context unavailable)
        """
        if self._in_transaction:
            raise RuntimeError("Nested transactions are not supported")

        # Check for out-of-transaction modifications
        current_state = msgspec.to_builtins(self.db)
        if current_state != self._previous_builtins:
            # Allow bootstrap to create a new database from empty state
            is_bootstrap = action in _BOOTSTRAP_ACTIONS
            if is_bootstrap and not self._previous_builtins:
                pass  # Expected: creating database from scratch
            else:
                diff = compute_diff(self._previous_builtins, current_state)
                diff_json = msgspec.json.encode(diff).decode()
                _logger.critical(
                    "Database state modified outside of transaction! "
                    "This indicates a bug where DB changes occurred without a transaction wrapper.\n"
                    f"Changes detected:\n{diff_json}"
                )
                raise SystemExit(1)

        old_action = self._current_action
        old_user = self._current_user
        self._current_action = action
        # Prefer ctx.user.uuid if ctx provided, otherwise use user param
        self._current_user = str(ctx.user.uuid) if ctx else user
        self._in_transaction = True
        self._transaction_snapshot = current_state

        try:
            yield
            current = msgspec.to_builtins(self.db)
            self._queue_change(
                self._current_action, self._current_version, current, self._current_user
            )
        except Exception:
            # Rollback on error: restore from snapshot
            _logger.warning("Transaction '%s' failed, rolling back changes", action)
            if self._transaction_snapshot is not None:
                decoder = msgspec.json.Decoder(DB)
                self.db = decoder.decode(
                    msgspec.json.encode(self._transaction_snapshot)
                )
                self.db._store = self
            raise
        finally:
            self._current_action = old_action
            self._current_user = old_user
            self._in_transaction = False
            self._transaction_snapshot = None

    async def flush(self) -> None:
        """Write all pending changes to disk."""
        await flush_changes(self.db_path, self._pending_changes)
