database.py

"""
database.py — SQLite Database Layer with SQLAlchemy
====================================================
Stores keys, invoices, credentials, and audit logs locally.
طبقة قاعدة البيانات لتخزين المفاتيح والفواتير والاعتمادات وسجلات التدقيق.

Models:
- KeyRecord         : EC private keys (encrypted at rest)
- InvoiceRecord     : Invoice XML, hash, QR, status
- CredentialRecord  : CSID tokens and secrets per ZATCA stage
- AuditLog          : API request/response audit trail
"""

from __future__ import annotations

import enum
from contextlib import contextmanager
from datetime import datetime
from typing import Generator, List, Optional, TypeVar

from sqlalchemy import (
    JSON,
    Column,
    DateTime,
    Enum,
    Float,
    Integer,
    LargeBinary,
    String,
    Text,
    create_engine,
    inspect,
)
from sqlalchemy.orm import (
    Session,
    declarative_base,
    sessionmaker,
)

from config import get_settings

# ---------------------------------------------------------------------------
# SQLAlchemy boilerplate / الإعداد الأساسي
# ---------------------------------------------------------------------------

settings = get_settings()
engine = create_engine(
    settings.database_url,
    connect_args={"check_same_thread": False} if settings.database_url.startswith("sqlite") else {},
    echo=settings.debug,
    future=True,
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


# ---------------------------------------------------------------------------
# Enum helpers / مساعدات الأنواع
# ---------------------------------------------------------------------------

class DBInvoiceStatus(str, enum.Enum):
    """Database-level invoice status values"""
    DRAFT = "draft"
    SIGNED = "signed"
    SUBMITTED = "submitted"
    CLEARED = "cleared"
    REJECTED = "rejected"
    CANCELLED = "cancelled"


class DBInvoiceType(str, enum.Enum):
    """Database-level invoice type values"""
    SIMPLIFIED = "simplified"
    STANDARD = "standard"
    CREDIT_NOTE = "credit_note"
    DEBIT_NOTE = "debit_note"


# ---------------------------------------------------------------------------
# Models / النماذج
# ---------------------------------------------------------------------------

class KeyRecord(Base):
    """
    Stores an EC private key encrypted at rest.
    يخزن المفتاح الخاص مشفراً
    """
    __tablename__ = "keys"

    id = Column(Integer, primary_key=True, autoincrement=True)
    key_id = Column(String(64), unique=True, nullable=False, index=True)
    private_key_pem = Column(Text, nullable=False)  # encrypted PEM
    public_key_pem = Column(Text, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)

    def __repr__(self) -> str:
        return f"<KeyRecord(id={self.id}, key_id='{self.key_id}')>"


class InvoiceRecord(Base):
    """
    Stores invoice data including XML payload, hash, QR code, and status.
    يخزن بيانات الفاتورة بما في ذلك XML والهاش وQR والحالة.
    """
    __tablename__ = "invoices"

    id = Column(Integer, primary_key=True, autoincrement=True)

    # Identification
    invoice_id = Column(String(64), unique=True, nullable=False, index=True)
    uuid_ = Column(String(36), unique=True, nullable=False, index=True)
    invoice_counter = Column(Integer, nullable=False)

    # Type & status
    invoice_type = Column(Enum(DBInvoiceType), nullable=False, default=DBInvoiceType.SIMPLIFIED)
    status = Column(Enum(DBInvoiceStatus), nullable=False, default=DBInvoiceStatus.DRAFT)

    # Merchant info (denormalized for quick retrieval)
    merchant_name = Column(String(255), nullable=False)
    merchant_vat = Column(String(15), nullable=False)

    # Totals
    total_net = Column(Float, nullable=False, default=0.0)
    total_vat = Column(Float, nullable=False, default=0.0)
    total_with_vat = Column(Float, nullable=False, default=0.0)

    # Payloads
    xml_payload = Column(Text, nullable=True)
    signed_xml = Column(Text, nullable=True)
    invoice_hash = Column(String(128), nullable=True)
    qr_code = Column(Text, nullable=True)  # base64-encoded QR image or TLV

    # Previous invoice hash (for chaining)
    previous_invoice_hash = Column(String(128), nullable=True)

    # For credit/debit notes
    original_invoice_id = Column(String(64), nullable=True, index=True)

    # Cancellation
    cancellation_reason = Column(Text, nullable=True)
    cancelled_by = Column(String(128), nullable=True)
    cancelled_at = Column(DateTime, nullable=True)

    # ZATCA response
    zatca_status = Column(String(32), nullable=True)
    zatca_response = Column(Text, nullable=True)

    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)

    def __repr__(self) -> str:
        return (
            f"<InvoiceRecord(id={self.id}, invoice_id='{self.invoice_id}', "
            f"type='{self.invoice_type.value}', status='{self.status.value}')>"
        )


class CredentialRecord(Base):
    """
    Stores ZATCA CSID credentials per stage (compliance / production).
    يخزن بيانات الاعتماد الخاصة بـ ZATCA لكل مرحلة.
    """
    __tablename__ = "credentials"

    id = Column(Integer, primary_key=True, autoincrement=True)
    credentials_id = Column(String(128), unique=True, nullable=False, index=True)
    binary_security_token = Column(Text, nullable=False)
    secret = Column(String(255), nullable=False)
    stage = Column(String(32), nullable=False)  # 'compliance' | 'production'
    request_id = Column(String(128), nullable=True)

    # Reference to the key used during onboarding
    key_id = Column(String(64), nullable=True)

    # Metadata
    is_active = Column(Integer, default=1, nullable=False)  # 1=active, 0=revoked
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    revoked_at = Column(DateTime, nullable=True)

    def __repr__(self) -> str:
        return (
            f"<CredentialRecord(id={self.id}, credentials_id='{self.credentials_id}', "
            f"stage='{self.stage}', active={self.is_active})>"
        )


class AuditLog(Base):
    """
    Audit trail for all API calls.
    سجل التدقيق لجميع استدعاءات API.
    """
    __tablename__ = "audit_log"

    id = Column(Integer, primary_key=True, autoincrement=True)

    # Request metadata
    request_id = Column(String(16), nullable=False, index=True)
    method = Column(String(10), nullable=False)
    endpoint = Column(String(255), nullable=False)
    path_params = Column(JSON, nullable=True)
    query_params = Column(JSON, nullable=True)
    request_body = Column(Text, nullable=True)

    # Response metadata
    status_code = Column(Integer, nullable=True)
    response_body = Column(Text, nullable=True)
    error_message = Column(Text, nullable=True)

    # Timing
    duration_ms = Column(Float, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)

    def __repr__(self) -> str:
        return (
            f"<AuditLog(id={self.id}, method='{self.method}', "
            f"endpoint='{self.endpoint}', status={self.status_code})>"
        )


# ---------------------------------------------------------------------------
# DB lifecycle / دورة حياة قاعدة البيانات
# ---------------------------------------------------------------------------

def init_db() -> None:
    """
    Create all tables if they don't exist.
    إنشاء الجداول إذا لم تكن موجودة.
    """
    Base.metadata.create_all(bind=engine)


def get_db() -> Generator[Session, None, None]:
    """
    FastAPI dependency: yields a database session and closes it.
    مولد الجلسات لـ FastAPI dependency injection.
    """
    session = SessionLocal()
    try:
        yield session
    finally:
        session.close()


@contextmanager
def db_session() -> Generator[Session, None, None]:
    """
    Context manager for manual DB sessions (non-FastAPI contexts).
    مدير سياق للجلسات اليدوية.
    """
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


# ---------------------------------------------------------------------------
# CRUD helpers / دوال الإضافة والقراءة والتحديث والحذف
# ---------------------------------------------------------------------------

T = TypeVar("T", bound=Base)


# -- Key CRUD --

def create_key(
    session: Session,
    key_id: str,
    private_key_pem: str,
    public_key_pem: Optional[str] = None,
) -> KeyRecord:
    """Insert a new key record — إدراج مفتاح جديد"""
    record = KeyRecord(
        key_id=key_id,
        private_key_pem=private_key_pem,
        public_key_pem=public_key_pem,
    )
    session.add(record)
    session.commit()
    session.refresh(record)
    return record


def get_key_by_id(session: Session, key_id: str) -> Optional[KeyRecord]:
    """Fetch key by its key_id — جلب مفتاح حسب المعرف"""
    return session.query(KeyRecord).filter(KeyRecord.key_id == key_id).first()


def list_keys(session: Session, skip: int = 0, limit: int = 100) -> List[KeyRecord]:
    """List keys with pagination — قائمة المفاتيح"""
    return session.query(KeyRecord).offset(skip).limit(limit).all()


# -- Invoice CRUD --

def create_invoice(
    session: Session,
    invoice_id: str,
    uuid_: str,
    invoice_counter: int,
    invoice_type: DBInvoiceType,
    merchant_name: str,
    merchant_vat: str,
    total_net: float,
    total_vat: float,
    total_with_vat: float,
    xml_payload: Optional[str] = None,
    signed_xml: Optional[str] = None,
    invoice_hash: Optional[str] = None,
    qr_code: Optional[str] = None,
    previous_invoice_hash: Optional[str] = None,
    original_invoice_id: Optional[str] = None,
    status: DBInvoiceStatus = DBInvoiceStatus.DRAFT,
) -> InvoiceRecord:
    """Insert a new invoice — إدراج فاتورة جديدة"""
    record = InvoiceRecord(
        invoice_id=invoice_id,
        uuid_=uuid_,
        invoice_counter=invoice_counter,
        invoice_type=invoice_type,
        status=status,
        merchant_name=merchant_name,
        merchant_vat=merchant_vat,
        total_net=total_net,
        total_vat=total_vat,
        total_with_vat=total_with_vat,
        xml_payload=xml_payload,
        signed_xml=signed_xml,
        invoice_hash=invoice_hash,
        qr_code=qr_code,
        previous_invoice_hash=previous_invoice_hash,
        original_invoice_id=original_invoice_id,
    )
    session.add(record)
    session.commit()
    session.refresh(record)
    return record


def get_invoice_by_id(session: Session, invoice_id: str) -> Optional[InvoiceRecord]:
    """Fetch invoice by invoice_id — جلب فاتورة حسب المعرف"""
    return session.query(InvoiceRecord).filter(InvoiceRecord.invoice_id == invoice_id).first()


def get_invoice_by_uuid(session: Session, uuid_: str) -> Optional[InvoiceRecord]:
    """Fetch invoice by UUID — جلب فاتورة حسب UUID"""
    return session.query(InvoiceRecord).filter(InvoiceRecord.uuid_ == uuid_).first()


def list_invoices(
    session: Session,
    status: Optional[DBInvoiceStatus] = None,
    invoice_type: Optional[DBInvoiceType] = None,
    skip: int = 0,
    limit: int = 100,
) -> List[InvoiceRecord]:
    """List invoices with optional filters — قائمة الفواتير"""
    query = session.query(InvoiceRecord)
    if status:
        query = query.filter(InvoiceRecord.status == status)
    if invoice_type:
        query = query.filter(InvoiceRecord.invoice_type == invoice_type)
    return query.order_by(InvoiceRecord.created_at.desc()).offset(skip).limit(limit).all()


def update_invoice_status(
    session: Session,
    invoice_id: str,
    new_status: DBInvoiceStatus,
    zatca_status: Optional[str] = None,
    zatca_response: Optional[str] = None,
) -> Optional[InvoiceRecord]:
    """Update invoice status — تحديث حالة الفاتورة"""
    record = get_invoice_by_id(session, invoice_id)
    if record is None:
        return None
    record.status = new_status
    if zatca_status:
        record.zatca_status = zatca_status
    if zatca_response:
        record.zatca_response = zatca_response
    record.updated_at = datetime.utcnow()
    session.commit()
    session.refresh(record)
    return record


def update_invoice_xml_and_hash(
    session: Session,
    invoice_id: str,
    signed_xml: Optional[str] = None,
    invoice_hash: Optional[str] = None,
    qr_code: Optional[str] = None,
) -> Optional[InvoiceRecord]:
    """Update signed XML, hash, and QR — تحديث XML الموقع والهاش وQR"""
    record = get_invoice_by_id(session, invoice_id)
    if record is None:
        return None
    if signed_xml:
        record.signed_xml = signed_xml
    if invoice_hash:
        record.invoice_hash = invoice_hash
    if qr_code:
        record.qr_code = qr_code
    record.updated_at = datetime.utcnow()
    session.commit()
    session.refresh(record)
    return record


def cancel_invoice(
    session: Session,
    invoice_id: str,
    reason: str,
    cancelled_by: str,
) -> Optional[InvoiceRecord]:
    """
    Cancel an invoice — إلغاء فاتورة.
    Only allowed if status is not CLEARED or already CANCELLED.
    """
    record = get_invoice_by_id(session, invoice_id)
    if record is None:
        return None
    if record.status in (DBInvoiceStatus.CLEARED, DBInvoiceStatus.CANCELLED):
        raise ValueError(
            f"لا يمكن إلغاء الفاتورة بحالة {record.status.value} | "
            f"Cannot cancel invoice in status {record.status.value}"
        )
    record.status = DBInvoiceStatus.CANCELLED
    record.cancellation_reason = reason
    record.cancelled_by = cancelled_by
    record.cancelled_at = datetime.utcnow()
    record.updated_at = datetime.utcnow()
    session.commit()
    session.refresh(record)
    return record


# -- Credential CRUD --

def create_credential(
    session: Session,
    credentials_id: str,
    binary_security_token: str,
    secret: str,
    stage: str,
    request_id: Optional[str] = None,
    key_id: Optional[str] = None,
) -> CredentialRecord:
    """Insert new CSID credentials — إدراج بيانات اعتماد جديدة"""
    record = CredentialRecord(
        credentials_id=credentials_id,
        binary_security_token=binary_security_token,
        secret=secret,
        stage=stage,
        request_id=request_id,
        key_id=key_id,
    )
    session.add(record)
    session.commit()
    session.refresh(record)
    return record


def get_credential_by_id(session: Session, credentials_id: str) -> Optional[CredentialRecord]:
    """Fetch credential by credentials_id — جلب الاعتماد"""
    return session.query(CredentialRecord).filter(
        CredentialRecord.credentials_id == credentials_id
    ).first()


def get_active_credentials(session: Session, stage: str) -> List[CredentialRecord]:
    """Get active credentials for a stage — جلب الاعتمادات النشطة"""
    return session.query(CredentialRecord).filter(
        CredentialRecord.stage == stage,
        CredentialRecord.is_active == 1,
    ).all()


def revoke_credential(session: Session, credentials_id: str) -> Optional[CredentialRecord]:
    """Revoke a credential — إبطال الاعتماد"""
    record = get_credential_by_id(session, credentials_id)
    if record is None:
        return None
    record.is_active = 0
    record.revoked_at = datetime.utcnow()
    session.commit()
    session.refresh(record)
    return record


# -- Audit Log CRUD --

def create_audit_log(
    session: Session,
    request_id: str,
    method: str,
    endpoint: str,
    path_params: Optional[dict] = None,
    query_params: Optional[dict] = None,
    request_body: Optional[str] = None,
    status_code: Optional[int] = None,
    response_body: Optional[str] = None,
    error_message: Optional[str] = None,
    duration_ms: Optional[float] = None,
) -> AuditLog:
    """Insert an audit log entry — إدراج سجل تدقيق"""
    record = AuditLog(
        request_id=request_id,
        method=method,
        endpoint=endpoint,
        path_params=path_params,
        query_params=query_params,
        request_body=request_body,
        status_code=status_code,
        response_body=response_body,
        error_message=error_message,
        duration_ms=duration_ms,
    )
    session.add(record)
    session.commit()
    session.refresh(record)
    return record


def list_audit_logs(
    session: Session,
    endpoint: Optional[str] = None,
    skip: int = 0,
    limit: int = 100,
) -> List[AuditLog]:
    """List audit logs with optional endpoint filter — قائمة سجلات التدقيق"""
    query = session.query(AuditLog)
    if endpoint:
        query = query.filter(AuditLog.endpoint == endpoint)
    return query.order_by(AuditLog.created_at.desc()).offset(skip).limit(limit).all()