from __future__ import annotations

from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Literal, Protocol

from sqlalchemy.orm import Session, sessionmaker


DatastoreMode = Literal["shared", "schema", "database"]


@dataclass(frozen=True, slots=True)
class TenantDatastoreRef:
    tenant_id: str
    mode: DatastoreMode = "shared"
    datastore_id: str | None = None
    schema_name: str | None = None
    database_url_secret_ref: str | None = None


class TenantDatastore(Protocol):
    ref: TenantDatastoreRef

    @contextmanager
    def session_for(self, module_id: str) -> Generator[Session, None, None]:
        ...


class DatastoreResolver(Protocol):
    def for_tenant(self, tenant_id: str) -> TenantDatastore:
        ...


class SharedTenantDatastore:
    def __init__(self, ref: TenantDatastoreRef, session_factory: sessionmaker[Session]) -> None:
        self.ref = ref
        self._session_factory = session_factory

    @contextmanager
    def session_for(self, module_id: str) -> Generator[Session, None, None]:
        del module_id
        with self._session_factory() as session:
            yield session


class SharedDatastoreResolver:
    def __init__(self, session_factory: sessionmaker[Session]) -> None:
        self._session_factory = session_factory

    def for_tenant(self, tenant_id: str) -> TenantDatastore:
        return SharedTenantDatastore(TenantDatastoreRef(tenant_id=tenant_id), self._session_factory)