db.py

"""
Database layer for persistence and audit trails.

Uses SQLite with SQLAlchemy for storing:
- User profiles
- Recommendation runs
- Scheme decisions
- Complete decision traces
"""

from sqlalchemy import create_engine, Column, String, Integer, Float, DateTime, JSON, Boolean, ForeignKey, Text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from datetime import datetime
from pathlib import Path
import json
from typing import Optional
import uuid

# Database setup
DATABASE_URL = "sqlite:///./policy_recommendations.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


class UserProfileDB(Base):
    """Persist user profiles for auditing."""
    __tablename__ = "user_profiles"
    
    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    age = Column(Integer, nullable=False)
    income = Column(Integer, nullable=False)
    state = Column(String(50), nullable=False)
    category = Column(String(20), nullable=False)
    gender = Column(String(20), nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow)
    
    # Relationships
    recommendation_runs = relationship("RecommendationRunDB", back_populates="user_profile")


class RecommendationRunDB(Base):
    """Persist each recommendation run with metadata."""
    __tablename__ = "recommendation_runs"
    
    run_id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    user_profile_id = Column(String, ForeignKey("user_profiles.id"), nullable=False)
    scoring_mode = Column(String(20), nullable=False)  # 'rules', 'ml', or 'hybrid'
    config_version = Column(String(50), default="1.0")
    created_at = Column(DateTime, default=datetime.utcnow)
    
    # Store summary metrics
    total_schemes_checked = Column(Integer, default=0)
    eligible_count = Column(Integer, default=0)
    ineligible_count = Column(Integer, default=0)
    
    # Relationships
    user_profile = relationship("UserProfileDB", back_populates="recommendation_runs")
    scheme_decisions = relationship("SchemeDecisionDB", back_populates="recommendation_run")
    
    def to_dict(self):
        """Convert to dictionary for API responses."""
        return {
            "run_id": self.run_id,
            "user_profile_id": self.user_profile_id,
            "scoring_mode": self.scoring_mode,
            "created_at": self.created_at.isoformat(),
            "summary": {
                "total_schemes": self.total_schemes_checked,
                "eligible": self.eligible_count,
                "ineligible": self.ineligible_count
            }
        }


class SchemeDecisionDB(Base):
    """Persist individual scheme eligibility and scoring decisions."""
    __tablename__ = "scheme_decisions"
    
    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    run_id = Column(String, ForeignKey("recommendation_runs.run_id"), nullable=False)
    scheme_id = Column(String, nullable=False)
    scheme_name = Column(String(255), nullable=False)
    
    # Decision details
    is_eligible = Column(Boolean, nullable=False)
    eligibility_reason = Column(Text)  # For ineligible: reason; for eligible: brief explanation
    
    # Scoring (if eligible)
    score = Column(Float, nullable=True)  # 0-100
    scoring_method = Column(String(20), nullable=True)  # 'rules', 'ml', or 'hybrid'
    
    # Full structured data
    decision_trace = Column(JSON)  # Complete DecisionTrace as JSON
    scoring_factors = Column(JSON)  # ScoringFactors as JSON
    ml_features = Column(JSON, nullable=True)  # ML feature contributions (if applicable)
    
    # ML Model tracking and confidence
    ml_model_version = Column(String(50), nullable=True)  # e.g. "logistic_v1.0" for ML/hybrid, null for rules-only
    ml_confidence = Column(Float, nullable=True)  # Predicted probability from ML model (0-1), null for rules-only
    
    created_at = Column(DateTime, default=datetime.utcnow)
    
    # Relationships
    recommendation_run = relationship("RecommendationRunDB", back_populates="scheme_decisions")
    
    def to_dict(self):
        """Convert to dictionary for API responses."""
        return {
            "scheme_id": self.scheme_id,
            "scheme_name": self.scheme_name,
            "is_eligible": self.is_eligible,
            "eligibility_reason": self.eligibility_reason,
            "score": self.score,
            "scoring_method": self.scoring_method,
            "decision_trace": self.decision_trace,
            "scoring_factors": self.scoring_factors,
            "ml_features": self.ml_features,
            "ml_model_version": self.ml_model_version,
            "ml_confidence": self.ml_confidence
        }


class DatabaseManager:
    """Singleton manager for database operations."""
    
    def __init__(self):
        self.SessionLocal = SessionLocal
    
    def init_db(self):
        """Initialize database tables."""
        Base.metadata.create_all(bind=engine)
    
    def get_session(self):
        """Get a database session."""
        return self.SessionLocal()
    
    def save_recommendation_run(self, user_profile: dict, scoring_mode: str, decisions: list) -> str:
        """
        Save a complete recommendation run.
        
        Args:
            user_profile: User data dict
            scoring_mode: Scoring method used ('rules', 'ml', 'hybrid')
            decisions: List of scheme decision dicts
        
        Returns:
            run_id for later audit lookup
        
        **Audit Immutability:** Once saved, audit records cannot be modified or deleted.
        """
        session = self.get_session()
        try:
            # Create or get user profile
            user_db = UserProfileDB(
                age=user_profile["age"],
                income=user_profile["income"],
                state=user_profile["state"],
                category=user_profile["category"],
                gender=user_profile["gender"]
            )
            session.add(user_db)
            session.flush()  # Get user ID
            
            # Create recommendation run
            eligible_count = sum(1 for d in decisions if d.get("is_eligible", False))
            ineligible_count = len(decisions) - eligible_count
            
            run_db = RecommendationRunDB(
                user_profile_id=user_db.id,
                scoring_mode=scoring_mode,
                total_schemes_checked=len(decisions),
                eligible_count=eligible_count,
                ineligible_count=ineligible_count
            )
            session.add(run_db)
            session.flush()  # Get run_id
            
            # Create scheme decisions
            for decision in decisions:
                scheme_decision = SchemeDecisionDB(
                    run_id=run_db.run_id,
                    scheme_id=decision["scheme_id"],
                    scheme_name=decision["scheme_name"],
                    is_eligible=decision.get("is_eligible", False),
                    eligibility_reason=decision.get("eligibility_reason"),
                    score=decision.get("score"),
                    scoring_method=decision.get("scoring_method"),
                    decision_trace=decision.get("decision_trace"),
                    scoring_factors=decision.get("scoring_factors"),
                    ml_features=decision.get("ml_features"),
                    ml_model_version=decision.get("ml_model_version"),
                    ml_confidence=decision.get("ml_confidence")
                )
                session.add(scheme_decision)
            
            session.commit()
            run_id = run_db.run_id
            return run_id
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()
    
    def get_audit_trail(self, run_id: str) -> dict:
        """
        Retrieve complete audit trail for a recommendation run.
        
        Args:
            run_id: The recommendation run ID
        
        Returns:
            Complete audit record with user, decisions, and traces
        """
        session = self.get_session()
        try:
            run = session.query(RecommendationRunDB).filter(
                RecommendationRunDB.run_id == run_id
            ).first()
            
            if not run:
                return None
            
            user = run.user_profile
            decisions = session.query(SchemeDecisionDB).filter(
                SchemeDecisionDB.run_id == run_id
            ).all()
            
            audit_record = {
                "run_id": run.run_id,
                "timestamp": run.created_at.isoformat(),
                "scoring_mode": run.scoring_mode,
                "user_profile": {
                    "age": user.age,
                    "income": user.income,
                    "state": user.state,
                    "category": user.category,
                    "gender": user.gender
                },
                "summary": {
                    "total_schemes_checked": run.total_schemes_checked,
                    "eligible_count": run.eligible_count,
                    "ineligible_count": run.ineligible_count
                },
                "decisions": [d.to_dict() for d in decisions]
            }
            
            return audit_record
        finally:
            session.close()
    
    def ensure_audit_immutability(self, run_id: str) -> None:
        """
        Verify audit record exists and enforce immutability.
        
        Audit records are write-once, read-many (WORM).
        This method verifies a record exists; attempts to modify it will fail
        at the application layer or database constraint layer.
        
        Args:
            run_id: The recommendation run ID
        
        Raises:
            ValueError: If run_id not found (immutable record not committed)
        """
        session = self.get_session()
        try:
            run = session.query(RecommendationRunDB).filter(
                RecommendationRunDB.run_id == run_id
            ).first()
            
            if not run:
                raise ValueError(f"Audit record not found: {run_id}")
        finally:
            session.close()


# Singleton instance
db_manager = DatabaseManager()


def get_db():
    """Get database manager instance."""
    return db_manager