app/database.py

import os
import json
import asyncpg
import structlog
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

logger = structlog.get_logger()

class PostcareDB:
    def __init__(self):
        self.pool = None
        self.db_url = os.getenv("DATABASE_URL")
        # Performance: Tuned Pool Settings
        self.min_pool = int(os.getenv("DB_MIN_POOL", "5"))
        self.max_pool = int(os.getenv("DB_MAX_POOL", "20"))

    async def connect(self):
        """Creates a connection pool to Neon."""
        if not self.pool:
            if not self.db_url:
                raise ValueError("DATABASE_URL environment variable is not set")
            self.pool = await asyncpg.create_pool(
                self.db_url,
                min_size=self.min_pool,
                max_size=self.max_pool,
                command_timeout=30,
                max_queries=50000
            )

    async def disconnect(self):
        """Closes the pool."""
        if self.pool:
            await self.pool.close()

    # Resilience: Retry DB queries on transient network errors
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=0.5, min=1, max=10),
        retry=retry_if_exception_type(OSError)
    )
    async def execute(self, query: str, params: list = None):
        """Executes INSERT/UPDATE/DELETE queries with retry logic."""
        if not self.pool:
            await self.connect()
        async with self.pool.acquire() as conn:
            return await conn.execute(query, *params if params else [])

    async def fetchrow(self, query: str, params: list = None):
        """Fetch a single row."""
        if not self.pool:
            await self.connect()
        async with self.pool.acquire() as conn:
            return await conn.fetchrow(query, *params if params else [])

    async def init_db(self):
        """Creates tables if they don't exist (Enterprise Schema)."""
        if not self.pool:
            await self.connect()
            
        async with self.pool.acquire() as conn:
            # 1. Hospitals (Tenants)
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS hospitals (
                    id TEXT PRIMARY KEY, 
                    name TEXT NOT NULL,
                    email TEXT UNIQUE NOT NULL,
                    tin TEXT NOT NULL, 
                    reg_no TEXT NOT NULL, 
                    mobile TEXT NOT NULL,
                    license_status TEXT DEFAULT 'active',
                    domain_verified BOOLEAN DEFAULT FALSE,
                    sso_enabled BOOLEAN DEFAULT FALSE,
                    directory_sync_enabled BOOLEAN DEFAULT FALSE,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)
            
            # 2. Patients
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS patients (
                    id TEXT PRIMARY KEY,
                    hospital_id TEXT,
                    full_name TEXT NOT NULL,
                    mobile TEXT NOT NULL,
                    dob TEXT NOT NULL, 
                    govt_id_last_four TEXT NOT NULL,
                    govt_id_full TEXT, 
                    gender TEXT,
                    next_of_kin_mobile TEXT,
                    medical_history_summary TEXT,
                    discharge_date TEXT,
                    symptoms_detected TEXT,
                    medications TEXT,
                    monitoring_status TEXT DEFAULT 'active',
                    FOREIGN KEY (hospital_id) REFERENCES hospitals(id)
                );
            """)

            # 3. Call Logs
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS call_logs (
                    id SERIAL PRIMARY KEY,
                    patient_id TEXT,
                    scheduled_time TIMESTAMP,
                    call_status TEXT DEFAULT 'pending',
                    agent_summary TEXT,
                    symptoms_detected TEXT,
                    adherence_confirmed BOOLEAN,
                    FOREIGN KEY (patient_id) REFERENCES patients(id)
                );
            """)

            # 4. Staff Users (Synced from WorkOS)
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS staff_users (
                    id TEXT PRIMARY KEY,
                    hospital_id TEXT NOT NULL,
                    email TEXT NOT NULL,
                    first_name TEXT,
                    last_name TEXT,
                    status TEXT DEFAULT 'active',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE (hospital_id, email),
                    FOREIGN KEY (hospital_id) REFERENCES hospitals(id)
                );
            """)

            # 5. Directory Groups (RBAC)
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS directory_groups (
                    id TEXT PRIMARY KEY,
                    hospital_id TEXT NOT NULL,
                    name TEXT NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (hospital_id) REFERENCES hospitals(id)
                );
            """)

            # 6. Group Membership
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS directory_group_memberships (
                    group_id TEXT NOT NULL,
                    user_id TEXT NOT NULL,
                    hospital_id TEXT NOT NULL,
                    PRIMARY KEY (group_id, user_id),
                    FOREIGN KEY (group_id) REFERENCES directory_groups(id),
                    FOREIGN KEY (user_id) REFERENCES staff_users(id),
                    FOREIGN KEY (hospital_id) REFERENCES hospitals(id)
                );
            """)

            # 7. Sessions (Stateful Auth)
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    hospital_id TEXT NOT NULL,
                    user_id TEXT NOT NULL,
                    refresh_token_hash TEXT NOT NULL,
                    user_agent TEXT,
                    ip TEXT,
                    revoked_at TIMESTAMP WITH TIME ZONE,
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                    last_used_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (hospital_id) REFERENCES hospitals(id),
                    FOREIGN KEY (user_id) REFERENCES staff_users(id)
                );
            """)

            # 8. Webhook Idempotency
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS webhook_events (
                    id TEXT PRIMARY KEY,
                    event_type TEXT NOT NULL,
                    received_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)

            # 9. Audit Logs
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS audit_logs (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    hospital_id TEXT NOT NULL,
                    user_id TEXT NOT NULL,
                    action TEXT NOT NULL,
                    entity_type TEXT,
                    entity_id TEXT,
                    changes JSONB,
                    ip_address TEXT,
                    user_agent TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)

            # Performance Indexes
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_patients_hospital_id ON patients(hospital_id);
                CREATE INDEX IF NOT EXISTS idx_call_logs_patient_id ON call_logs(patient_id);
                CREATE INDEX IF NOT EXISTS idx_call_logs_scheduled_time ON call_logs(scheduled_time);
                CREATE INDEX IF NOT EXISTS idx_sessions_active ON sessions(id) WHERE revoked_at IS NULL;
                CREATE INDEX IF NOT EXISTS idx_audit_hospital_time ON audit_logs(hospital_id, created_at);
            """)

db = PostcareDB()