from __future__ import annotations

import logging
from pathlib import Path
from typing import TypedDict

from diskcache import Index

from solutions.graph.models import GraphDBReadInProgressError, GraphDBWriteInProgressError

logger = logging.getLogger(__name__)


class _DbAccessState(TypedDict):
    db_path: Path | str | None
    read_handles_count: int
    write_locked: bool


class DbReadWriteLock:
    """
    Manages safe concurrent access to a DbInstance via disk-backed access state.

    Ensures:
      - Multiple concurrent readers allowed.
      - Single writer exclusive access.
      - Prevents runtime conflicts in multi-process environments.
    """

    def __init__(self, access_index: Index) -> None:
        self.access_index = access_index

    def try_acquire_lock(self, db_id: Path | str, readonly: bool) -> None:
        """
        Raises:
            GraphDBWriteInProgressError
            GraphDBReadInProgressError
        """
        with self.access_index.transact():
            state: _DbAccessState | None = self.access_index.get(db_id)

            if state:
                if state["write_locked"]:
                    logger.warning("Database %s is already write locked, cannot read or write. State: %s", db_id, state)
                    # As long as it is write locked, it is not usable.
                    raise GraphDBWriteInProgressError()

                if not readonly and state["read_handles_count"] > 0:
                    logger.warning("Database %s is already read locked, cannot write. State: %s", db_id, state)
                    # If we want a write lock but there are reads happening already, we cannot write the graph.
                    raise GraphDBReadInProgressError()

                state = _DbAccessState(**state)
                if readonly:
                    state["read_handles_count"] += 1
                else:
                    state["write_locked"] = True
            else:
                # First access
                state = _DbAccessState(
                    db_path=db_id,
                    read_handles_count=1 if readonly else 0,
                    write_locked=not readonly,
                )

            self.access_index[db_id] = state
            logger.debug(f"Acquired db lock: state for %s is now %s.", db_id, state)

    def release_lock(self, db_id: Path | str, readonly: bool) -> None:
        """
        Raises:
            RuntimeError
        """
        with self.access_index.transact():
            state: _DbAccessState | None = self.access_index.get(db_id)
            if not state:
                raise RuntimeError("Should not happen")

            state = _DbAccessState(**state)
            if readonly:
                if state["read_handles_count"] == 0:
                    logger.error(
                        "Unexpected db lock state. Tried to release a read lock when none were held. State for %s is: %s",
                        db_id,
                        state,
                    )
                    raise RuntimeError("Tried to release a read lock when none were held.")

                state["read_handles_count"] = max(0, state["read_handles_count"] - 1)
            else:
                state["write_locked"] = False

            self.access_index[db_id] = state
            logger.debug("Released db lock: state for %s is now %s.", db_id, state)
