ljwx-stock/backend/app/dao/recommendation_snapshot_dao.py

"""
Data Access Object for RecommendationSnapshot
Provides database operations for versioned recommendation snapshots
"""

from typing import Optional, List, Dict, Any
from datetime import datetime, date
from sqlalchemy import desc, and_
from sqlalchemy.orm import Session
from sqlalchemy.exc import IntegrityError

from app.models.recommendation_batch_models import RecommendationSnapshot
from app.database import SessionLocal


class RecommendationSnapshotDAO:
    """
    DAO for managing recommendation snapshots in PostgreSQL

    Key Methods:
    - create_snapshot: Create a new draft snapshot
    - publish_snapshot: Atomically publish a draft snapshot
    - get_latest_published: Get the latest published snapshot
    - get_by_version: Get a specific version
    - list_snapshots: List snapshots with filters
    - archive_old_snapshots: Archive snapshots older than retention period
    """

    def __init__(self, db: Optional[Session] = None):
        """
        Initialize DAO with optional database session
        If no session provided, creates a new one
        """
        self.db = db or SessionLocal()
        self._own_session = db is None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._own_session:
            self.db.close()

    def create_snapshot(
        self,
        version: str,
        model_version: str,
        data_date: date,
        result_json: Dict[str, Any],
        market_context: Optional[Dict[str, Any]] = None,
        explanation: Optional[str] = None,
        stock_universe: Optional[List[str]] = None,
        feature_summary: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> Optional[RecommendationSnapshot]:
        """
        Create a new recommendation snapshot in 'draft' status

        Args:
            version: Unique version string (e.g., "20251225_0330")
            model_version: Model version used (e.g., "v1.2.0")
            data_date: Data cutoff date
            result_json: Complete recommendation result
            market_context: Optional market environment data
            explanation: Optional LLM-generated explanation
            stock_universe: Optional list of candidate stocks
            feature_summary: Optional feature metadata
            **kwargs: Additional fields (total_stocks, recommended_count, avg_score, etc.)

        Returns:
            Created RecommendationSnapshot or None on error
        """
        try:
            snapshot = RecommendationSnapshot(
                version=version,
                model_version=model_version,
                data_date=data_date,
                generate_time=datetime.now(),
                result_json=result_json,
                market_context=market_context,
                explanation=explanation,
                stock_universe=stock_universe,
                feature_summary=feature_summary,
                status='draft',
                **kwargs
            )

            self.db.add(snapshot)
            self.db.commit()
            self.db.refresh(snapshot)

            return snapshot

        except IntegrityError as e:
            self.db.rollback()
            print(f"❌ Failed to create snapshot (version conflict): {e}")
            return None
        except Exception as e:
            self.db.rollback()
            print(f"❌ Failed to create snapshot: {e}")
            return None

    def publish_snapshot(self, version: str) -> bool:
        """
        Atomically publish a draft snapshot
        Sets publish_time and updates status to 'published'

        Args:
            version: Version to publish

        Returns:
            True if successful, False otherwise
        """
        try:
            snapshot = self.db.query(RecommendationSnapshot).filter(
                and_(
                    RecommendationSnapshot.version == version,
                    RecommendationSnapshot.status == 'draft'
                )
            ).first()

            if not snapshot:
                print(f"❌ Snapshot {version} not found or not in draft status")
                return False

            snapshot.status = 'published'
            snapshot.publish_time = datetime.now()

            self.db.commit()
            print(f"✅ Published snapshot {version}")
            return True

        except Exception as e:
            self.db.rollback()
            print(f"❌ Failed to publish snapshot: {e}")
            return False

    def get_latest_published(self) -> Optional[RecommendationSnapshot]:
        """
        Get the latest published recommendation snapshot
        Ordered by publish_time DESC

        Returns:
            Latest published snapshot or None
        """
        return self.db.query(RecommendationSnapshot).filter(
            RecommendationSnapshot.status == 'published'
        ).order_by(
            desc(RecommendationSnapshot.publish_time)
        ).first()

    def get_by_version(self, version: str) -> Optional[RecommendationSnapshot]:
        """
        Get a specific snapshot by version

        Args:
            version: Version string

        Returns:
            Snapshot or None
        """
        return self.db.query(RecommendationSnapshot).filter(
            RecommendationSnapshot.version == version
        ).first()

    def list_snapshots(
        self,
        status: Optional[str] = None,
        limit: int = 10,
        offset: int = 0
    ) -> List[RecommendationSnapshot]:
        """
        List snapshots with optional filters

        Args:
            status: Filter by status (draft/published/archived)
            limit: Max number of results
            offset: Pagination offset

        Returns:
            List of snapshots
        """
        query = self.db.query(RecommendationSnapshot)

        if status:
            query = query.filter(RecommendationSnapshot.status == status)

        return query.order_by(
            desc(RecommendationSnapshot.generate_time)
        ).limit(limit).offset(offset).all()

    def archive_snapshot(self, version: str) -> bool:
        """
        Archive a snapshot (change status to 'archived')

        Args:
            version: Version to archive

        Returns:
            True if successful
        """
        try:
            snapshot = self.db.query(RecommendationSnapshot).filter(
                RecommendationSnapshot.version == version
            ).first()

            if not snapshot:
                return False

            snapshot.status = 'archived'
            self.db.commit()
            return True

        except Exception as e:
            self.db.rollback()
            print(f"❌ Failed to archive snapshot: {e}")
            return False

    def delete_snapshot(self, version: str) -> bool:
        """
        Permanently delete a snapshot
        Use with caution - prefer archive_snapshot for safety

        Args:
            version: Version to delete

        Returns:
            True if successful
        """
        try:
            snapshot = self.db.query(RecommendationSnapshot).filter(
                RecommendationSnapshot.version == version
            ).first()

            if not snapshot:
                return False

            self.db.delete(snapshot)
            self.db.commit()
            return True

        except Exception as e:
            self.db.rollback()
            print(f"❌ Failed to delete snapshot: {e}")
            return False

    def get_snapshot_count(self, status: Optional[str] = None) -> int:
        """
        Get count of snapshots

        Args:
            status: Optional status filter

        Returns:
            Count of snapshots
        """
        query = self.db.query(RecommendationSnapshot)

        if status:
            query = query.filter(RecommendationSnapshot.status == status)

        return query.count()

    def get_snapshots_by_date_range(
        self,
        start_date: date,
        end_date: date,
        status: Optional[str] = None
    ) -> List[RecommendationSnapshot]:
        """
        Get snapshots within a date range

        Args:
            start_date: Start date (inclusive)
            end_date: End date (inclusive)
            status: Optional status filter

        Returns:
            List of snapshots
        """
        query = self.db.query(RecommendationSnapshot).filter(
            and_(
                RecommendationSnapshot.data_date >= start_date,
                RecommendationSnapshot.data_date <= end_date
            )
        )

        if status:
            query = query.filter(RecommendationSnapshot.status == status)

        return query.order_by(
            desc(RecommendationSnapshot.data_date)
        ).all()

    def update_snapshot_error(self, version: str, error_message: str) -> bool:
        """
        Update snapshot with error information
        Useful for marking failed generation attempts

        Args:
            version: Version to update
            error_message: Error description

        Returns:
            True if successful
        """
        try:
            snapshot = self.db.query(RecommendationSnapshot).filter(
                RecommendationSnapshot.version == version
            ).first()

            if not snapshot:
                return False

            snapshot.error_message = error_message
            snapshot.status = 'draft'  # Keep as draft on error

            self.db.commit()
            return True

        except Exception as e:
            self.db.rollback()
            print(f"❌ Failed to update snapshot error: {e}")
            return False