database.py

"""
Database setup and models for ZenVolt asset management.
Uses SQLAlchemy with SQLite for simplicity.
"""

from sqlalchemy import create_engine, Column, String, Float, Integer, DateTime, Text, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from datetime import datetime
import json
from typing import Optional, List, Dict

import os

# Database URL (SQLite for local dev, PostgreSQL for Production)
DATABASE_URL = os.getenv("DATABASE_URL")

# Detect Container environment (Docker/Hugging Face)
IS_DOCKER = os.path.exists("/.dockerenv") or os.environ.get("HOSTNAME", "").endswith("-space")

if not DATABASE_URL:
    if IS_DOCKER:
        # Check if we have a persistent volume mount
        if os.path.exists("/data"):
            DATABASE_URL = "sqlite:////data/zenvolt_assets.db"
        else:
            # Fallback to /tmp for writable storage in ephemeral containers
            DATABASE_URL = "sqlite:////tmp/zenvolt_assets.db"
    else:
        # Local development
        DATABASE_URL = "sqlite:///./zenvolt_assets.db"

# SQLite specific config
is_sqlite = DATABASE_URL.startswith("sqlite")
connect_args = {"check_same_thread": False} if is_sqlite else {}

engine = create_engine(DATABASE_URL, connect_args=connect_args)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


# Database Models
class Asset(Base):
    """Battery Asset model"""
    __tablename__ = "assets"
    
    id = Column(String, primary_key=True, index=True)
    owner_id = Column(String, index=True)  # User who owns this asset
    project_name = Column(String, nullable=False)
    asset_id = Column(String, unique=True, nullable=False, index=True)  # e.g., "BESS_01"
    location = Column(String)
    capacity = Column(Float)  # kWh
    bms_provider = Column(String)
    category = Column(String)  # JSON array as string: ["Grid-Scale", "Certified"]
    project_stage = Column(String)  # Certified, Monitoring, Under Review, Pending
    
    # Health metrics (latest)
    health_score = Column(Float)  # 0-100
    remaining_life = Column(Float)  # years
    degradation_rate = Column(Float)  # % per year
    confidence = Column(Float)  # 0-1
    
    # Operational data
    state_of_charge = Column(Float)  # percentage
    voltage = Column(Float)  # V
    temperature = Column(Float)  # °C
    is_online = Column(Integer, default=1)  # Boolean as int
    
    # NFT/Blockchain (optional)
    nft_hash = Column(String)
    nft_mint_date = Column(DateTime)
    block_explorer_url = Column(String)
    
    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    last_prediction_date = Column(DateTime)
    
    # Relationships
    predictions = relationship("HealthPrediction", back_populates="asset", cascade="all, delete-orphan")
    
    def to_dict(self) -> Dict:
        """Convert to dictionary for API response"""
        return {
            "id": self.id,
            "owner_id": self.owner_id,
            "projectName": self.project_name,
            "assetId": self.asset_id,
            "location": self.location,
            "capacity": self.capacity,
            "bmsProvider": self.bms_provider,
            "category": json.loads(self.category) if self.category else [],
            "projectStage": self.project_stage,
            "healthScore": self.health_score,
            "remainingLife": self.remaining_life,
            "degradationRate": self.degradation_rate,
            "confidence": self.confidence,
            "stateOfCharge": self.state_of_charge,
            "voltage": self.voltage,
            "temperature": self.temperature,
            "isOnline": bool(self.is_online),
            "nftHash": self.nft_hash,
            "nftMintDate": self.nft_mint_date.isoformat() if self.nft_mint_date else None,
            "blockExplorerUrl": self.block_explorer_url,
            "createdAt": self.created_at.isoformat() if self.created_at else None,
            "updatedAt": self.updated_at.isoformat() if self.updated_at else None,
            "lastUpdated": self.last_prediction_date.isoformat() if self.last_prediction_date else None,
        }


class HealthPrediction(Base):
    """Historical health predictions for assets"""
    __tablename__ = "health_predictions"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    asset_id = Column(String, ForeignKey("assets.id"), nullable=False, index=True)
    
    # Prediction timestamp
    prediction_date = Column(DateTime, default=datetime.utcnow, nullable=False, index=True)
    cycle_index = Column(Integer)  # Cycle number at time of prediction
    
    # Health metrics
    health_score = Column(Float, nullable=False)  # 0-100
    current_soh = Column(Float)  # 0-1
    soh_lower_95 = Column(Float)  # Confidence interval
    soh_upper_95 = Column(Float)
    
    remaining_life_years = Column(Float)
    remaining_life_cycles = Column(Float)
    degradation_rate = Column(Float)  # % per year
    
    confidence = Column(Float)  # 0-1
    data_quality = Column(Float)  # 0-1
    
    # SHAP explanation (top features as JSON)
    shap_top_features = Column(Text)  # JSON string
    
    # Relationship
    asset = relationship("Asset", back_populates="predictions")
    
    def to_dict(self) -> Dict:
        """Convert to dictionary for API response"""
        return {
            "id": self.id,
            "asset_id": self.asset_id,
            "prediction_date": self.prediction_date.isoformat() if self.prediction_date else None,
            "cycle_index": self.cycle_index,
            "health_score": self.health_score,
            "current_soh": self.current_soh,
            "soh_lower_95": self.soh_lower_95,
            "soh_upper_95": self.soh_upper_95,
            "remaining_life_years": self.remaining_life_years,
            "remaining_life_cycles": self.remaining_life_cycles,
            "degradation_rate": self.degradation_rate,
            "confidence": self.confidence,
            "data_quality": self.data_quality,
            "shap_top_features": json.loads(self.shap_top_features) if self.shap_top_features else None,
        }


# Database session dependency for FastAPI
def get_db():
    """Dependency to get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


# Initialize database
def init_db():
    """Create all tables"""
    Base.metadata.create_all(bind=engine)
    print("✓ Database initialized")


if __name__ == "__main__":
    init_db()