trading-journal/backend/trading_journal/crud.py

from __future__ import annotations

from datetime import date, datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, TypeVar, cast

from pydantic import BaseModel
from sqlalchemy.exc import IntegrityError
from sqlmodel import Session, select

from trading_journal import models

if TYPE_CHECKING:
    from collections.abc import Mapping
    from enum import Enum

    from sqlalchemy.sql.elements import ColumnElement


# Generic enum member type
T = TypeVar("T", bound="Enum")


def _check_enum(enum_cls: type[T], value: object, field_name: str) -> T:
    if value is None:
        raise ValueError(f"{field_name} is required")
    # already an enum member
    if isinstance(value, enum_cls):
        return value
    # strict string match: must match exactly enum name or enum value (case-sensitive)
    if isinstance(value, str):
        for m in enum_cls:
            if m.name == value or str(m.value) == value:
                return m
    allowed = [m.name for m in enum_cls]
    raise ValueError(f"Invalid {field_name!s}: {value!r}. Allowed: {allowed}")


def _allowed_columns(model: type[models.SQLModel]) -> set[str]:
    tbl = cast("models.SQLModel", model).__table__  # type: ignore[attr-defined]
    return {c.name for c in tbl.columns}


AnyModel = Any


def _data_to_dict(data: AnyModel) -> dict[str, AnyModel]:
    if isinstance(data, BaseModel):
        return data.model_dump(exclude_unset=True)
    if hasattr(data, "dict"):
        return data.dict(exclude_unset=True)
    return dict(data)


# Trades
def create_trade(session: Session, trade_data: Mapping[str, Any] | BaseModel) -> models.Trades:
    data = _data_to_dict(trade_data)
    allowed = _allowed_columns(models.Trades)
    payload = {k: v for k, v in data.items() if k in allowed}
    cycle_id = payload.get("cycle_id")
    if "symbol" not in payload:
        raise ValueError("symbol is required")
    if "exchange_id" not in payload and cycle_id is None:
        raise ValueError("exchange_id is required when no cycle is attached")
    # If an exchange_id is provided (and no cycle is attached), ensure the exchange exists
    # and belongs to the same user as the trade (if user_id is provided).
    if cycle_id is None and "exchange_id" in payload:
        ex = session.get(models.Exchanges, payload["exchange_id"])
        if ex is None:
            raise ValueError("exchange_id does not exist")
        user_id = payload.get("user_id")
        if user_id is not None and ex.user_id != user_id:
            raise ValueError("exchange.user_id does not match trade.user_id")
    if "underlying_currency" not in payload:
        raise ValueError("underlying_currency is required")
    payload["underlying_currency"] = _check_enum(models.UnderlyingCurrency, payload["underlying_currency"], "underlying_currency")
    if "trade_type" not in payload:
        raise ValueError("trade_type is required")
    payload["trade_type"] = _check_enum(models.TradeType, payload["trade_type"], "trade_type")
    if "trade_strategy" not in payload:
        raise ValueError("trade_strategy is required")
    payload["trade_strategy"] = _check_enum(models.TradeStrategy, payload["trade_strategy"], "trade_strategy")
    # trade_time_utc is the creation moment: always set to now (caller shouldn't provide)
    now = datetime.now(timezone.utc)
    payload.pop("trade_time_utc", None)
    payload["trade_time_utc"] = now
    if "trade_date" not in payload or payload.get("trade_date") is None:
        payload["trade_date"] = payload["trade_time_utc"].date()
    user_id = payload.get("user_id")
    if "quantity" not in payload:
        raise ValueError("quantity is required")
    if "price_cents" not in payload:
        raise ValueError("price_cents is required")
    if "commission_cents" not in payload:
        payload["commission_cents"] = 0
    if "gross_cash_flow_cents" not in payload:
        raise ValueError("gross_cash_flow_cents is required")
    if "net_cash_flow_cents" not in payload:
        raise ValueError("net_cash_flow_cents is required")

    # If no cycle_id provided, create Cycle instance but don't call create_cycle()
    created_cycle = None
    if cycle_id is None:
        c_payload = {
            "user_id": user_id,
            "symbol": payload["symbol"],
            "exchange_id": payload["exchange_id"],
            "underlying_currency": payload["underlying_currency"],
            "friendly_name": "Auto-created Cycle by trade " + payload.get("friendly_name", ""),
            "status": models.CycleStatus.OPEN,
            "start_date": payload["trade_date"],
        }
        created_cycle = models.Cycles(**c_payload)
        session.add(created_cycle)
        # do NOT flush here; will flush together with trade below

    # If cycle_id provided, validate existence and ownership
    if cycle_id is not None:
        cycle = session.get(models.Cycles, cycle_id)

        if cycle is None:
            raise ValueError("cycle_id does not exist")
        payload.pop("exchange_id", None)  # ignore exchange_id if provided; use cycle's exchange_id
        payload["exchange_id"] = cycle.exchange_id
        if cycle.user_id != user_id:
            raise ValueError("cycle.user_id does not match trade.user_id")

    # Build trade instance; if we created a Cycle instance, link via relationship so a single flush will persist both and populate ids
    t_payload = dict(payload)
    # remove cycle_id if we're using created_cycle; relationship will set it on flush
    if created_cycle is not None:
        t_payload.pop("cycle_id", None)
    t = models.Trades(**t_payload)
    if created_cycle is not None:
        t.cycle = created_cycle

    session.add(t)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("create_trade integrity error") from e
    session.refresh(t)
    return t


def get_trade_by_id(session: Session, trade_id: int) -> models.Trades | None:
    return session.get(models.Trades, trade_id)


def get_trade_by_user_id_and_friendly_name(session: Session, user_id: int, friendly_name: str) -> models.Trades | None:
    statement = select(models.Trades).where(
        models.Trades.user_id == user_id,
        models.Trades.friendly_name == friendly_name,
    )
    return session.exec(statement).first()


def get_trades_by_user_id(session: Session, user_id: int) -> list[models.Trades]:
    statement = select(models.Trades).where(
        models.Trades.user_id == user_id,
    )
    return list(session.exec(statement).all())


def update_trade_friendly_name(session: Session, trade_id: int, friendly_name: str) -> models.Trades:
    trade: models.Trades | None = session.get(models.Trades, trade_id)
    if trade is None:
        raise ValueError("trade_id does not exist")
    trade.friendly_name = friendly_name
    session.add(trade)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("update_trade_friendly_name integrity error") from e
    session.refresh(trade)
    return trade


def update_trade_note(session: Session, trade_id: int, note: str) -> models.Trades:
    trade: models.Trades | None = session.get(models.Trades, trade_id)
    if trade is None:
        raise ValueError("trade_id does not exist")
    trade.notes = note
    session.add(trade)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("update_trade_note integrity error") from e
    session.refresh(trade)
    return trade


def invalidate_trade(session: Session, trade_id: int) -> models.Trades:
    trade: models.Trades | None = session.get(models.Trades, trade_id)
    if trade is None:
        raise ValueError("trade_id does not exist")
    if trade.is_invalidated:
        raise ValueError("trade is already invalidated")
    trade.is_invalidated = True
    trade.invalidated_at = datetime.now(timezone.utc)
    session.add(trade)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("invalidate_trade integrity error") from e
    session.refresh(trade)
    return trade


def replace_trade(session: Session, old_trade_id: int, new_trade_data: Mapping[str, Any] | BaseModel) -> models.Trades:
    invalidate_trade(session, old_trade_id)
    data = _data_to_dict(new_trade_data)
    data["replaced_by_trade_id"] = old_trade_id
    return create_trade(session, data)


# Cycles
def create_cycle(session: Session, cycle_data: Mapping[str, Any] | BaseModel) -> models.Cycles:
    data = _data_to_dict(cycle_data)
    allowed = _allowed_columns(models.Cycles)
    payload = {k: v for k, v in data.items() if k in allowed}
    if "user_id" not in payload:
        raise ValueError("user_id is required")
    if "symbol" not in payload:
        raise ValueError("symbol is required")
    if "exchange_id" not in payload:
        raise ValueError("exchange_id is required")
    # ensure the exchange exists and belongs to the same user
    ex = session.get(models.Exchanges, payload["exchange_id"])
    if ex is None:
        raise ValueError("exchange_id does not exist")
    if ex.user_id != payload.get("user_id"):
        raise ValueError("exchange.user_id does not match cycle.user_id")
    if "underlying_currency" not in payload:
        raise ValueError("underlying_currency is required")
    payload["underlying_currency"] = _check_enum(models.UnderlyingCurrency, payload["underlying_currency"], "underlying_currency")
    if "status" not in payload:
        raise ValueError("status is required")
    payload["status"] = _check_enum(models.CycleStatus, payload["status"], "status")
    if "start_date" not in payload:
        raise ValueError("start_date is required")

    c = models.Cycles(**payload)
    session.add(c)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("create_cycle integrity error") from e
    session.refresh(c)
    return c


IMMUTABLE_CYCLE_FIELDS = {"id", "user_id", "start_date"}


def get_cycle_by_id(session: Session, cycle_id: int) -> models.Cycles | None:
    return session.get(models.Cycles, cycle_id)


def get_cycles_by_user_id(session: Session, user_id: int) -> list[models.Cycles]:
    statement = select(models.Cycles).where(
        models.Cycles.user_id == user_id,
    )
    return list(session.exec(statement).all())


def update_cycle(session: Session, cycle_id: int, update_data: Mapping[str, Any] | BaseModel) -> models.Cycles:
    cycle: models.Cycles | None = session.get(models.Cycles, cycle_id)
    if cycle is None:
        raise ValueError("cycle_id does not exist")
    data = _data_to_dict(update_data)

    allowed = _allowed_columns(models.Cycles)
    for k, v in data.items():
        if k in IMMUTABLE_CYCLE_FIELDS:
            raise ValueError(f"field {k!r} is immutable")
        if k not in allowed:
            continue
        # If trying to change exchange_id, ensure the new exchange exists and belongs to
        # the same user as the cycle.
        if k == "exchange_id":
            ex = session.get(models.Exchanges, v)
            if ex is None:
                raise ValueError("exchange_id does not exist")
            if ex.user_id != cycle.user_id:
                raise ValueError("exchange.user_id does not match cycle.user_id")
        if k == "underlying_currency":
            v = _check_enum(models.UnderlyingCurrency, v, "underlying_currency")  # noqa: PLW2901
        if k == "status":
            v = _check_enum(models.CycleStatus, v, "status")  # noqa: PLW2901
        setattr(cycle, k, v)
    session.add(cycle)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("update_cycle integrity error") from e
    session.refresh(cycle)
    return cycle


# Cycle loan and interest
def create_cycle_loan_event(session: Session, loan_data: Mapping[str, Any] | BaseModel) -> models.CycleLoanChangeEvents:
    data = _data_to_dict(loan_data)
    allowed = _allowed_columns(models.CycleLoanChangeEvents)
    payload = {k: v for k, v in data.items() if k in allowed}
    if "cycle_id" not in payload:
        raise ValueError("cycle_id is required")
    cycle = session.get(models.Cycles, payload["cycle_id"])
    if cycle is None:
        raise ValueError("cycle_id does not exist")

    payload["effective_date"] = payload.get("effective_date") or datetime.now(timezone.utc).date()
    payload["created_at"] = datetime.now(timezone.utc)
    cle = models.CycleLoanChangeEvents(**payload)
    session.add(cle)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("create_cycle_loan_event integrity error") from e
    session.refresh(cle)
    return cle


def get_loan_events_by_cycle_id(session: Session, cycle_id: int) -> list[models.CycleLoanChangeEvents]:
    eff_col = cast("ColumnElement", models.CycleLoanChangeEvents.effective_date)
    id_col = cast("ColumnElement", models.CycleLoanChangeEvents.id)
    statement = (
        select(models.CycleLoanChangeEvents)
        .where(
            models.CycleLoanChangeEvents.cycle_id == cycle_id,
        )
        .order_by(eff_col, id_col.asc())
    )
    return list(session.exec(statement).all())


def get_loan_event_by_cycle_id_and_effective_date(session: Session, cycle_id: int, effective_date: date) -> models.CycleLoanChangeEvents | None:
    statement = select(models.CycleLoanChangeEvents).where(
        models.CycleLoanChangeEvents.cycle_id == cycle_id,
        models.CycleLoanChangeEvents.effective_date == effective_date,
    )
    return session.exec(statement).first()


def update_cycle_loan_event(session: Session, event_id: int, update_data: Mapping[str, Any] | BaseModel) -> models.CycleLoanChangeEvents:
    event: models.CycleLoanChangeEvents | None = session.get(models.CycleLoanChangeEvents, event_id)
    if event is None:
        raise ValueError("event_id does not exist")
    data = _data_to_dict(update_data)

    allowed = _allowed_columns(models.CycleLoanChangeEvents)
    for k, v in data.items():
        if k in {"id", "cycle_id", "effective_date", "created_at"}:
            raise ValueError(f"field {k!r} is immutable")
        if k not in allowed:
            continue
        setattr(event, k, v)
    session.add(event)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("update_cycle_loan_event integrity error") from e
    session.refresh(event)
    return event


def create_cycle_daily_accrual(session: Session, cycle_id: int, accrual_date: date, accrual_amount_cents: int) -> models.CycleDailyAccrual:
    cycle = session.get(models.Cycles, cycle_id)
    if cycle is None:
        raise ValueError("cycle_id does not exist")
    existing = session.exec(
        select(models.CycleDailyAccrual).where(
            models.CycleDailyAccrual.cycle_id == cycle_id,
            models.CycleDailyAccrual.accrual_date == accrual_date,
        ),
    ).first()
    if existing:
        return existing
    if accrual_amount_cents < 0:
        raise ValueError("accrual_amount_cents must be non-negative")
    row = models.CycleDailyAccrual(
        cycle_id=cycle_id,
        accrual_date=accrual_date,
        accrual_amount_cents=accrual_amount_cents,
        created_at=datetime.now(timezone.utc),
    )
    session.add(row)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("create_cycle_daily_accrual integrity error") from e
    session.refresh(row)
    return row


def get_cycle_daily_accruals_by_cycle_id(session: Session, cycle_id: int) -> list[models.CycleDailyAccrual]:
    date_col = cast("ColumnElement", models.CycleDailyAccrual.accrual_date)
    statement = (
        select(models.CycleDailyAccrual)
        .where(
            models.CycleDailyAccrual.cycle_id == cycle_id,
        )
        .order_by(date_col.asc())
    )
    return list(session.exec(statement).all())


def get_cycle_daily_accrual_by_cycle_id_and_date(session: Session, cycle_id: int, accrual_date: date) -> models.CycleDailyAccrual | None:
    statement = select(models.CycleDailyAccrual).where(
        models.CycleDailyAccrual.cycle_id == cycle_id,
        models.CycleDailyAccrual.accrual_date == accrual_date,
    )
    return session.exec(statement).first()


# Exchanges
IMMUTABLE_EXCHANGE_FIELDS = {"id"}


def create_exchange(session: Session, exchange_data: Mapping[str, Any] | BaseModel) -> models.Exchanges:
    data = _data_to_dict(exchange_data)
    allowed = _allowed_columns(models.Exchanges)
    payload = {k: v for k, v in data.items() if k in allowed}
    if "name" not in payload:
        raise ValueError("name is required")

    e = models.Exchanges(**payload)
    session.add(e)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("create_exchange integrity error") from e
    session.refresh(e)
    return e


def get_exchange_by_id(session: Session, exchange_id: int) -> models.Exchanges | None:
    return session.get(models.Exchanges, exchange_id)


def get_exchange_by_name_and_user_id(session: Session, name: str, user_id: int) -> models.Exchanges | None:
    statement = select(models.Exchanges).where(
        models.Exchanges.name == name,
        models.Exchanges.user_id == user_id,
    )
    return session.exec(statement).first()


def get_all_exchanges(session: Session) -> list[models.Exchanges]:
    statement = select(models.Exchanges)
    return list(session.exec(statement).all())


def get_all_exchanges_by_user_id(session: Session, user_id: int) -> list[models.Exchanges]:
    statement = select(models.Exchanges).where(
        models.Exchanges.user_id == user_id,
    )
    return list(session.exec(statement).all())


def update_exchange(session: Session, exchange_id: int, update_data: Mapping[str, Any] | BaseModel) -> models.Exchanges:
    exchange: models.Exchanges | None = session.get(models.Exchanges, exchange_id)
    if exchange is None:
        raise ValueError("exchange_id does not exist")
    data = _data_to_dict(update_data)
    allowed = _allowed_columns(models.Exchanges)
    for k, v in data.items():
        if k in IMMUTABLE_EXCHANGE_FIELDS:
            raise ValueError(f"field {k!r} is immutable")
        if k in allowed:
            setattr(exchange, k, v)
    session.add(exchange)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("update_exchange integrity error") from e
    session.refresh(exchange)
    return exchange


def delete_exchange(session: Session, exchange_id: int) -> None:
    exchange: models.Exchanges | None = session.get(models.Exchanges, exchange_id)
    if exchange is None:
        return
    session.delete(exchange)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("delete_exchange integrity error") from e


# Users
IMMUTABLE_USER_FIELDS = {"id", "username", "created_at"}


def create_user(session: Session, user_data: Mapping[str, Any] | BaseModel) -> models.Users:
    data = _data_to_dict(user_data)
    allowed = _allowed_columns(models.Users)
    payload = {k: v for k, v in data.items() if k in allowed}
    if "username" not in payload:
        raise ValueError("username is required")
    if "password_hash" not in payload:
        raise ValueError("password_hash is required")

    u = models.Users(**payload)
    session.add(u)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("create_user integrity error") from e
    session.refresh(u)
    return u


def get_user_by_id(session: Session, user_id: int) -> models.Users | None:
    return session.get(models.Users, user_id)


def get_user_by_username(session: Session, username: str) -> models.Users | None:
    statement = select(models.Users).where(
        models.Users.username == username,
    )
    return session.exec(statement).first()


def update_user(session: Session, user_id: int, update_data: Mapping[str, Any] | BaseModel) -> models.Users:
    user: models.Users | None = session.get(models.Users, user_id)
    if user is None:
        raise ValueError("user_id does not exist")
    data = _data_to_dict(update_data)
    allowed = _allowed_columns(models.Users)
    for k, v in data.items():
        if k in IMMUTABLE_USER_FIELDS:
            raise ValueError(f"field {k!r} is immutable")
        if k in allowed:
            setattr(user, k, v)
    session.add(user)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("update_user integrity error") from e
    session.refresh(user)
    return user


# Sessions
def create_login_session(
    session: Session,
    user_id: int,
    session_token_hash: str,
    session_length_seconds: int = 86400,
    last_used_ip: str | None = None,
    user_agent: str | None = None,
    device_name: str | None = None,
) -> models.Sessions:
    user: models.Users | None = session.get(models.Users, user_id)
    if user is None:
        raise ValueError("user_id does not exist")
    user_id_val = cast("int", user.id)
    now = datetime.now(timezone.utc)
    expires_at = now + timedelta(seconds=session_length_seconds)
    s = models.Sessions(
        user_id=user_id_val,
        session_token_hash=session_token_hash,
        created_at=now,
        expires_at=expires_at,
        last_seen_at=now,
        last_used_ip=last_used_ip,
        user_agent=user_agent,
        device_name=device_name,
    )
    session.add(s)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("create_login_session integrity error") from e
    session.refresh(s)
    return s


def get_login_session_by_token_hash_and_user_id(session: Session, session_token_hash: str, user_id: int) -> models.Sessions | None:
    statement = select(models.Sessions).where(
        models.Sessions.session_token_hash == session_token_hash,
        models.Sessions.user_id == user_id,
        models.Sessions.expires_at > datetime.now(timezone.utc),
    )

    return session.exec(statement).first()


def get_login_session_by_token_hash(session: Session, session_token_hash: str) -> models.Sessions | None:
    statement = select(models.Sessions).where(
        models.Sessions.session_token_hash == session_token_hash,
        models.Sessions.expires_at > datetime.now(timezone.utc),
    )

    return session.exec(statement).first()


IMMUTABLE_SESSION_FIELDS = {"id", "user_id", "session_token_hash", "created_at"}


def update_login_session(
    session: Session,
    session_token_hashed: str,
    update_session: Mapping[str, Any] | BaseModel,
) -> models.Sessions | None:
    login_session: models.Sessions | None = session.exec(
        select(models.Sessions).where(
            models.Sessions.session_token_hash == session_token_hashed,
            models.Sessions.expires_at > datetime.now(timezone.utc),
        ),
    ).first()
    if login_session is None:
        return None
    data = _data_to_dict(update_session)
    allowed = _allowed_columns(models.Sessions)
    for k, v in data.items():
        if k in allowed and k not in IMMUTABLE_SESSION_FIELDS:
            setattr(login_session, k, v)
    session.add(login_session)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("update_login_session integrity error") from e
    session.refresh(login_session)
    return login_session


def delete_login_session(session: Session, session_token_hash: str) -> None:
    login_session: models.Sessions | None = session.exec(
        select(models.Sessions).where(
            models.Sessions.session_token_hash == session_token_hash,
        ),
    ).first()
    if login_session is None:
        return
    session.delete(login_session)
    try:
        session.flush()
    except IntegrityError as e:
        session.rollback()
        raise ValueError("delete_login_session integrity error") from e