"""Automatic per-portfolio filtering of the tables holding personal data.

Every `select()` on a session carrying a portfolio id gets a
`WHERE portfolio_id = :pid` on the models in `SCOPED_MODELS` — top-level
queries, eager loads and lazy loads alike.

It fails closed: a select touching a scoped model with no portfolio id on the
session raises. `unscoped_session()` is the only opt-out, and is greppable.

Writes are not covered — SQLAlchemy applies loader criteria to selects only —
so inserts set `portfolio_id` via `owned()`, with a before_flush net in
`install()`.

This separates users *within the application*, not from anyone holding the
database: see decisions.md #63 (Trust boundary).
"""

from __future__ import annotations

from contextlib import contextmanager
from typing import Iterator

from sqlalchemy import event, select
from sqlalchemy import false as sql_false
from sqlalchemy.orm import Session, with_loader_criteria

from .models import (
    MEMBER_ROLES,
    CorporateAction,
    Dividend,
    HoldingPref,
    InvestmentPlan,
    PlannedPurchase,
    PortfolioMember,
    SavedChart,
    Trade,
)

# Every model with a portfolio_id column. Adding a personal-data table? Add it
# here — that is the whole registration step.
SCOPED_MODELS = (
    Trade,
    Dividend,
    CorporateAction,
    PlannedPurchase,
    HoldingPref,
    InvestmentPlan,
)

# Owned by a person rather than a portfolio (their charts follow them between
# portfolios). Filtered the same way, on user_id — being in this tuple is what
# keeps them inside the fail-closed guarantee rather than relying on every
# query remembering a where clause.
USER_SCOPED_MODELS = (SavedChart,)

_PORTFOLIO_KEY = "portfolio_id"
_USER_KEY = "portfolio_user_id"
_BYPASS_KEY = "portfolio_unscoped"


class TenancyError(RuntimeError):
    """Raised when personal data is queried with no portfolio context."""


_UNBOUND_MESSAGE = (
    "query touches portfolio data with no portfolio on the session — use "
    "portfolio_session(...) or, for market-data maintenance, unscoped_session(...)"
)


def install(session_factory) -> None:
    """Attach the filter to a sessionmaker. Call once at startup."""

    @event.listens_for(session_factory, "before_flush")
    def _stamp_owner(session, _ctx, _instances) -> None:
        """New personal rows inherit the session's portfolio, and record who
        entered them.

        `owned()` says it explicitly at the call sites that matter; this is the
        safety net for bulk paths (the workbook importer creates rows in eight
        places) and for anything added later. Only fills blanks — it never
        reassigns an existing owner.
        """
        portfolio_id = session.info.get(_PORTFOLIO_KEY)
        user_id = session.info.get(_USER_KEY)
        if portfolio_id is None:
            return
        for obj in session.new:
            if not isinstance(obj, SCOPED_MODELS):
                continue
            if obj.portfolio_id is None:
                obj.portfolio_id = portfolio_id
            if user_id is not None and getattr(obj, "user_id", "absent") is None:
                obj.user_id = user_id

    @event.listens_for(session_factory, "do_orm_execute")
    def _apply_portfolio_filter(state) -> None:
        if not state.is_select or state.is_column_load:
            return
        session = state.session
        if session.info.get(_BYPASS_KEY):
            return
        portfolio_id = session.info.get(_PORTFOLIO_KEY)
        # A relationship load inherits the criteria from the query that loaded
        # its parent (propagate_to_loaders), so re-applying them here would
        # double up. But an UNBOUND session has no criteria to inherit, and a
        # lazy `instrument.trades` would then hand back every portfolio's rows
        # — so the refusal below still has to run before we return.
        if state.is_relationship_load:
            if portfolio_id is None and _touches_scoped(state):
                raise TenancyError(_UNBOUND_MESSAGE)
            return
        if portfolio_id is None:
            # No portfolio bound: fine for market data (the price feed runs this
            # way), a bug for anything personal.
            if _touches_scoped(state):
                raise TenancyError(_UNBOUND_MESSAGE)
            return
        # Applied to EVERY statement, not just ones naming a scoped model:
        # `select(Instrument)` carries the criteria forward to a later
        # `instrument.trades` lazy load (propagate_to_loaders). Skipping it here
        # because the statement "doesn't touch trades" is exactly how that lazy
        # load would come back unfiltered.
        for model in SCOPED_MODELS:
            state.statement = state.statement.options(
                with_loader_criteria(
                    model, model.portfolio_id == portfolio_id, include_aliases=True
                )
            )
        # User-scoped models are filtered whether or not a user is bound. With
        # no user, the criterion matches nothing rather than everything: a
        # session carrying a portfolio but no user (an API key whose creator
        # was deleted, nulling `created_by`) would otherwise read EVERY user's
        # saved charts.
        user_id = session.info.get(_USER_KEY)
        for model in USER_SCOPED_MODELS:
            criterion = (
                model.user_id == user_id if user_id is not None else sql_false()
            )
            state.statement = state.statement.options(
                with_loader_criteria(model, criterion, include_aliases=True)
            )


def _touches_scoped(state) -> bool:
    """True when the statement selects, or eager-loads, a scoped model.

    Cheap over-approximation: match on the mapper of every entity in the
    statement. Instrument-only queries (price feed, catalogue) skip the check so
    they keep working outside a request.
    """
    scoped = set(SCOPED_MODELS) | set(USER_SCOPED_MODELS)
    for desc in state.statement.column_descriptions:
        if desc.get("entity") in scoped:
            return True
    # A scoped table reached through FROM or JOIN rather than named as the
    # selected entity: `select(func.count()).select_from(Trade)` selects a
    # number, not a Trade, and `select(Instrument).join(Trade)` selects an
    # instrument — both read personal rows and neither shows up above.
    # Compared by table, because that is what a JOIN carries.
    scoped_tables = {model.__table__ for model in scoped}
    for element in getattr(state.statement, "get_final_froms", lambda: ())():
        if _mentions(element, scoped_tables):
            return True
    # Eager loads name their target in the statement's _with_options; simplest
    # reliable signal is the ORM path they walk, every step of it, so a scoped
    # model in the middle of a nested load counts too.
    for opt in getattr(state.statement, "_with_options", ()):
        for attr in getattr(opt, "path", ()) or ():
            if getattr(getattr(attr, "entity", None), "class_", None) in scoped:
                return True
    return False


def _mentions(element, tables: set) -> bool:
    """Whether a FROM element is, or joins, one of these tables."""
    if element in tables:
        return True
    left, right = getattr(element, "left", None), getattr(element, "right", None)
    if left is not None and _mentions(left, tables):
        return True
    if right is not None and _mentions(right, tables):
        return True
    # A subquery keeps its own SELECT, which may reach a scoped table too.
    inner = getattr(getattr(element, "element", None), "get_final_froms", None)
    if inner is not None:
        return any(_mentions(f, tables) for f in inner())
    return False


def bind(session: Session, portfolio_id: int, user_id: int | None = None) -> Session:
    """Confine an already-open session to one portfolio's rows.

    `user_id` is stamped onto new rows as "recorded by" — it plays no part in
    what the session can read.
    """
    session.info[_PORTFOLIO_KEY] = portfolio_id
    session.info[_USER_KEY] = user_id
    return session


def allow_unscoped(session: Session) -> Session:
    """Drop the filter on an open session. Two callers only: the first-run
    claim of pre-auth rows (which have no owner yet, so no filter could find
    them) and the price feed's cross-portfolio FX repair."""
    session.info[_BYPASS_KEY] = True
    return session


@contextmanager
def as_portfolio(
    session: Session,
    portfolio_id: int,
    *,
    user_id: int | None,
    roles: tuple[str, ...] = MEMBER_ROLES,
) -> Iterator[Session]:
    """Read another of THIS USER's portfolios for the length of the block.

    The one operation that spans two portfolios is moving a holding between
    them, and it must validate the destination's timeline before writing there.
    `unscoped_session` is the wrong tool for it: that drops the filter for
    every portfolio at once and says never to serve a request with it.

    The membership check is HERE, not in the caller. A primitive that rebinds
    the filter on request is a cross-tenant read waiting for one careless call
    site — so it refuses unless `user_id` holds a role in `roles`, and callers
    that are about to WRITE pass `roles=WRITE_ROLES` rather than checking
    separately and hoping the two agree.

    Reads inside the block must use `select()`. `session.get()` can answer from
    the identity map without going near the filter, so it would hand back a row
    belonging to whichever portfolio loaded it first.
    """
    if user_id is None:
        raise TenancyError("cannot reach another portfolio with no user")
    role = session.scalar(
        select(PortfolioMember.role).where(
            PortfolioMember.portfolio_id == portfolio_id,
            PortfolioMember.user_id == user_id,
        )
    )
    if role is None or role not in roles:
        # One message for "not a member" and "not allowed to write", because
        # telling them apart would confirm the portfolio exists.
        raise TenancyError(
            f"user {user_id} may not reach portfolio {portfolio_id} as {roles}"
        )
    was = session.info.get(_PORTFOLIO_KEY)
    session.info[_PORTFOLIO_KEY] = portfolio_id
    try:
        yield session
    finally:
        # Restored even when the block raises: a failed validation must not
        # leave the rest of the request reading the wrong portfolio.
        session.info[_PORTFOLIO_KEY] = was


@contextmanager
def portfolio_session(
    session_factory, portfolio_id: int, user_id: int | None = None
) -> Iterator[Session]:
    """A session whose reads are confined to one portfolio. Commits on success."""
    session = session_factory()
    session.info[_PORTFOLIO_KEY] = portfolio_id
    session.info[_USER_KEY] = user_id
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager
def unscoped_session(session_factory) -> Iterator[Session]:
    """Every portfolio's rows — market-data maintenance only (the price feed's
    trade-FX backfill). Never use this to serve a request."""
    session = session_factory()
    session.info[_BYPASS_KEY] = True
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def current_portfolio_id(session: Session) -> int | None:
    return session.info.get(_PORTFOLIO_KEY)


def current_user_id(session: Session) -> int | None:
    return session.info.get(_USER_KEY)


def owned(session: Session, obj):
    """Stamp a new row with the session's portfolio (and recorder) before add."""
    from .models import SavedChart as _SavedChart

    if isinstance(obj, _SavedChart):
        user_id = current_user_id(session)
        if user_id is None:
            raise TenancyError("cannot create a chart with no user context")
        obj.user_id = user_id
        obj.portfolio_id = current_portfolio_id(session)
        return obj
    portfolio_id = current_portfolio_id(session)
    if portfolio_id is None:
        raise TenancyError(
            f"cannot create {type(obj).__name__} with no portfolio context"
        )
    obj.portfolio_id = portfolio_id
    user_id = current_user_id(session)
    if user_id is not None and hasattr(obj, "user_id"):
        obj.user_id = user_id
    return obj