app/db/database.py

from datetime import datetime
from typing import AsyncGenerator

from sqlalchemy import Boolean, DateTime, Float, Integer, String, Text, event, DDL, CheckConstraint
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

from app.config import settings


class Base(DeclarativeBase):
    pass


# ------------------------------------------------------------------ #
#  Engine — asyncpg for PostgreSQL, aiosqlite fallback for tests      #
# ------------------------------------------------------------------ #
def _make_engine(url: str):
    is_sqlite = url.startswith("sqlite")
    return create_async_engine(
        url,
        echo=False,
        pool_pre_ping=True,  # detect stale connections
        **(
            {}
            if is_sqlite
            else {
                "pool_size": 10,
                "max_overflow": 20,
                "pool_timeout": 30,
            }
        ),
    )


engine = _make_engine(settings.DATABASE_URL)

AsyncSessionLocal = async_sessionmaker(
    bind=engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False,
)


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


# ------------------------------------------------------------------ #
#  ORM Models                                                          #
# ------------------------------------------------------------------ #


class DBUser(Base):
    __tablename__ = "users"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
    username: Mapped[str] = mapped_column(String(100), unique=True, index=True, nullable=False)
    email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
    password_hash: Mapped[str] = mapped_column(String(255), nullable=False)
    role: Mapped[str] = mapped_column(
        String(50),
        default="PASSENGER",
        # Valid roles enforced at DB level
    )
    zone: Mapped[str] = mapped_column(String(10), nullable=True)
    is_active: Mapped[bool] = mapped_column(Boolean, default=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)

    __table_args__ = (
        CheckConstraint(
            "role IN ('PASSENGER','STATION_AGENT','CONTROLLER','ADMIN')",
            name="ck_users_role",
        ),
    )


class DBStation(Base):
    __tablename__ = "stations"

    code: Mapped[str] = mapped_column(String(10), primary_key=True, index=True)
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    zone: Mapped[str] = mapped_column(String(10), nullable=False)
    division: Mapped[str] = mapped_column(String(50), nullable=True)
    latitude: Mapped[float] = mapped_column(Float, nullable=True)
    longitude: Mapped[float] = mapped_column(Float, nullable=True)
    is_major_junction: Mapped[bool] = mapped_column(Boolean, default=False)
    platform_count: Mapped[int] = mapped_column(Integer, nullable=True)


class DBSection(Base):
    __tablename__ = "sections"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
    from_station: Mapped[str] = mapped_column(String(10), nullable=False)
    to_station: Mapped[str] = mapped_column(String(10), nullable=False)
    distance_km: Mapped[float] = mapped_column(Float, nullable=False)
    max_speed_kmh: Mapped[int] = mapped_column(Integer, default=110)
    signaling_type: Mapped[str] = mapped_column(String(50), default="ABSOLUTE_BLOCK")
    capacity_trains_per_hour: Mapped[int] = mapped_column(Integer, default=10)


class DBDisruption(Base):
    __tablename__ = "disruptions"

    id: Mapped[str] = mapped_column(String(50), primary_key=True, index=True)
    train_no: Mapped[str] = mapped_column(String(10), nullable=False)
    section_from: Mapped[str] = mapped_column(String(10), nullable=False)
    section_to: Mapped[str] = mapped_column(String(10), nullable=False)
    disruption_type: Mapped[str] = mapped_column(String(50), default="DELAY_CASCADE")
    severity: Mapped[str] = mapped_column(String(20), default="MEDIUM")
    cascade_depth: Mapped[int] = mapped_column(Integer, default=0)
    trains_affected_json: Mapped[str] = mapped_column(Text, default="[]")
    passengers_affected: Mapped[int] = mapped_column(Integer, default=0)
    status: Mapped[str] = mapped_column(String(20), default="ACTIVE")
    detected_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    resolved_at: Mapped[datetime] = mapped_column(DateTime, nullable=True)

    __table_args__ = (
        CheckConstraint(
            "severity IN ('LOW','MEDIUM','HIGH','CRITICAL')",
            name="ck_disruptions_severity",
        ),
        CheckConstraint(
            "status IN ('ACTIVE','RESOLVED','ESCALATED')",
            name="ck_disruptions_status",
        ),
    )


class DBRecommendation(Base):
    __tablename__ = "recommendations"

    id: Mapped[str] = mapped_column(String(50), primary_key=True, index=True)
    disruption_id: Mapped[str] = mapped_column(String(50), nullable=False)
    type: Mapped[str] = mapped_column(String(20), default="HOLD")
    target_train: Mapped[str] = mapped_column(String(10), nullable=False)
    target_section: Mapped[str] = mapped_column(String(100), nullable=False)
    reasoning: Mapped[str] = mapped_column(Text, nullable=False)
    confidence: Mapped[float] = mapped_column(Float, nullable=False)
    tier: Mapped[int] = mapped_column(Integer, default=1)
    is_approved: Mapped[bool] = mapped_column(Boolean, default=False)
    override_reason: Mapped[str] = mapped_column(Text, nullable=True)
    generated_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)


class DBAuditEntry(Base):
    """
    APPEND-ONLY table.
    Enforcement layers:
      1. CheckConstraint prevents NULL hashes
      2. SQLAlchemy never exposes UPDATE/DELETE on this model
      3. PostgreSQL row-level security (applied via DDL below after table creation)
      4. Separate INSERT-only DB role in production Supabase setup
    """

    __tablename__ = "audit_log"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
    agent_name: Mapped[str] = mapped_column(String(100), nullable=False)
    action_type: Mapped[str] = mapped_column(String(100), nullable=False)
    target: Mapped[str] = mapped_column(String(200), nullable=False)
    reasoning: Mapped[str] = mapped_column(Text, nullable=False)
    confidence: Mapped[float] = mapped_column(Float, nullable=False)
    timestamp: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, index=True)
    prev_hash: Mapped[str] = mapped_column(String(64), nullable=False)
    current_hash: Mapped[str] = mapped_column(String(64), nullable=False, unique=True)

    __table_args__ = (
        CheckConstraint("length(current_hash) = 64", name="ck_audit_hash_length"),
        CheckConstraint("length(prev_hash) = 64", name="ck_audit_prev_hash_length"),
    )


# ------------------------------------------------------------------ #
#  PostgreSQL-only DDL: RLS append-only policy on audit_log           #
#  (Skipped for SQLite test environments)                             #
# ------------------------------------------------------------------ #
_AUDIT_RLS_DDL = DDL("""
    ALTER TABLE audit_log ENABLE ROW LEVEL SECURITY;
    DO $$
    BEGIN
        IF NOT EXISTS (
            SELECT 1 FROM pg_policies
            WHERE tablename = 'audit_log' AND policyname = 'audit_insert_only'
        ) THEN
            CREATE POLICY audit_insert_only ON audit_log FOR INSERT WITH CHECK (true);
        END IF;
    END
    $$;
""")

event.listen(
    DBAuditEntry.__table__,
    "after_create",
    _AUDIT_RLS_DDL.execute_if(dialect="postgresql"),
)


# ------------------------------------------------------------------ #
#  TimescaleDB hypertable for train_positions (applied post-create)   #
# ------------------------------------------------------------------ #
class DBTrainPosition(Base):
    """
    Time-series table. On PostgreSQL + TimescaleDB, converted to a
    hypertable via the DDL listener below.
    """

    __tablename__ = "train_positions"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
    train_no: Mapped[str] = mapped_column(String(10), nullable=False, index=True)
    train_name: Mapped[str] = mapped_column(String(255), nullable=True)
    at_station: Mapped[str] = mapped_column(String(10), nullable=True)
    scheduled_arrival: Mapped[datetime] = mapped_column(DateTime, nullable=True)
    actual_arrival: Mapped[datetime] = mapped_column(DateTime, nullable=True)
    delay_minutes: Mapped[int] = mapped_column(Integer, default=0)
    data_source: Mapped[str] = mapped_column(String(50), default="indianrailapi")
    data_quality: Mapped[float] = mapped_column(Float, default=1.0)
    recorded_at: Mapped[datetime] = mapped_column(
        DateTime, default=datetime.utcnow, index=True, nullable=False
    )


_TIMESCALE_DDL = DDL("""
    SELECT create_hypertable(
        'train_positions', 'recorded_at',
        if_not_exists => TRUE,
        migrate_data   => TRUE
    );
""")

event.listen(
    DBTrainPosition.__table__,
    "after_create",
    _TIMESCALE_DDL.execute_if(dialect="postgresql"),
)


# ------------------------------------------------------------------ #
#  Init                                                                #
# ------------------------------------------------------------------ #
async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)