from __future__ import annotations

from contextlib import contextmanager
from contextvars import ContextVar, Token
from dataclasses import dataclass
from typing import Iterator

from psycopg.conninfo import make_conninfo

from app.core.config import get_pgconn_string, get_timescaledb_pgconn_string


@dataclass(frozen=True)
class ActiveProjectRouting:
    project_code: str
    business_dsn: str
    timescale_dsn: str | None = None


_active_project_routing: ContextVar[ActiveProjectRouting | None] = ContextVar(
    "active_project_routing",
    default=None,
)


def get_active_project_routing() -> ActiveProjectRouting | None:
    return _active_project_routing.get()


@contextmanager
def activate_project_routing(
    routing: ActiveProjectRouting,
) -> Iterator[ActiveProjectRouting]:
    token: Token[ActiveProjectRouting | None] = _active_project_routing.set(routing)
    try:
        yield routing
    finally:
        _active_project_routing.reset(token)


def _dsn_for_database(dsn: str, database_name: str) -> str:
    return make_conninfo(dsn, dbname=database_name)


def get_project_pgconn_string(db_name: str | None = None) -> str:
    routing = get_active_project_routing()
    if routing is None:
        return get_pgconn_string(db_name=db_name)
    if db_name is None or db_name == routing.project_code:
        return routing.business_dsn
    return _dsn_for_database(routing.business_dsn, db_name)


def get_project_timescale_pgconn_string(db_name: str | None = None) -> str:
    routing = get_active_project_routing()
    if routing is None:
        return get_timescaledb_pgconn_string(db_name=db_name)
    if routing.timescale_dsn is None:
        raise RuntimeError(
            f"TimescaleDB routing is not configured for project {routing.project_code}"
        )
    # Legacy simulation code used to derive the Timescale database name from
    # the project code. Project-scoped requests must instead use the complete
    # iot_data DSN selected by metadata routing.
    return routing.timescale_dsn