"""
Prediction Batch Repository
Repository for prediction batch operations
"""

from typing import Optional, List, Dict, Any
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import text
from datetime import datetime, timedelta, timezone
import structlog

from .base import ForecastingBaseRepository
from app.models.forecasts import PredictionBatch
from shared.database.exceptions import DatabaseError, ValidationError

logger = structlog.get_logger()


class PredictionBatchRepository(ForecastingBaseRepository):
    """Repository for prediction batch operations"""
    
    def __init__(self, session: AsyncSession, cache_ttl: Optional[int] = 300):
        # Batch operations change frequently, shorter cache time (5 minutes)
        super().__init__(PredictionBatch, session, cache_ttl)
    
    async def create_batch(self, batch_data: Dict[str, Any]) -> PredictionBatch:
        """Create a new prediction batch"""
        try:
            # Validate batch data
            validation_result = self._validate_forecast_data(
                batch_data, 
                ["tenant_id", "batch_name"]
            )
            
            if not validation_result["is_valid"]:
                raise ValidationError(f"Invalid batch data: {validation_result['errors']}")
            
            # Set default values
            if "status" not in batch_data:
                batch_data["status"] = "pending"
            if "forecast_days" not in batch_data:
                batch_data["forecast_days"] = 7
            if "business_type" not in batch_data:
                batch_data["business_type"] = "individual"
            
            batch = await self.create(batch_data)
            
            logger.info("Prediction batch created",
                       batch_id=batch.id,
                       tenant_id=batch.tenant_id,
                       batch_name=batch.batch_name)
            
            return batch
            
        except ValidationError:
            raise
        except Exception as e:
            logger.error("Failed to create prediction batch",
                        tenant_id=batch_data.get("tenant_id"),
                        error=str(e))
            raise DatabaseError(f"Failed to create batch: {str(e)}")
    
    async def update_batch_progress(
        self,
        batch_id: str,
        completed_products: int = None,
        failed_products: int = None,
        total_products: int = None,
        status: str = None
    ) -> Optional[PredictionBatch]:
        """Update batch progress"""
        try:
            update_data = {}
            
            if completed_products is not None:
                update_data["completed_products"] = completed_products
            if failed_products is not None:
                update_data["failed_products"] = failed_products
            if total_products is not None:
                update_data["total_products"] = total_products
            if status:
                update_data["status"] = status
                if status in ["completed", "failed"]:
                    update_data["completed_at"] = datetime.now(timezone.utc)
            
            if not update_data:
                return await self.get_by_id(batch_id)
            
            updated_batch = await self.update(batch_id, update_data)
            
            logger.debug("Batch progress updated",
                        batch_id=batch_id,
                        status=status,
                        completed=completed_products)
            
            return updated_batch
            
        except Exception as e:
            logger.error("Failed to update batch progress",
                        batch_id=batch_id,
                        error=str(e))
            raise DatabaseError(f"Failed to update batch: {str(e)}")
    
    async def complete_batch(
        self,
        batch_id: str,
        processing_time_ms: int = None
    ) -> Optional[PredictionBatch]:
        """Mark batch as completed"""
        try:
            update_data = {
                "status": "completed",
                "completed_at": datetime.now(timezone.utc)
            }
            
            if processing_time_ms:
                update_data["processing_time_ms"] = processing_time_ms
            
            updated_batch = await self.update(batch_id, update_data)
            
            logger.info("Batch completed",
                       batch_id=batch_id,
                       processing_time_ms=processing_time_ms)
            
            return updated_batch
            
        except Exception as e:
            logger.error("Failed to complete batch",
                        batch_id=batch_id,
                        error=str(e))
            raise DatabaseError(f"Failed to complete batch: {str(e)}")
    
    async def fail_batch(
        self,
        batch_id: str,
        error_message: str,
        processing_time_ms: int = None
    ) -> Optional[PredictionBatch]:
        """Mark batch as failed"""
        try:
            update_data = {
                "status": "failed",
                "completed_at": datetime.now(timezone.utc),
                "error_message": error_message
            }
            
            if processing_time_ms:
                update_data["processing_time_ms"] = processing_time_ms
            
            updated_batch = await self.update(batch_id, update_data)
            
            logger.error("Batch failed",
                        batch_id=batch_id,
                        error_message=error_message)
            
            return updated_batch
            
        except Exception as e:
            logger.error("Failed to mark batch as failed",
                        batch_id=batch_id,
                        error=str(e))
            raise DatabaseError(f"Failed to fail batch: {str(e)}")
    
    async def cancel_batch(
        self,
        batch_id: str,
        cancelled_by: str = None
    ) -> Optional[PredictionBatch]:
        """Cancel a batch"""
        try:
            batch = await self.get_by_id(batch_id)
            if not batch:
                return None
            
            if batch.status in ["completed", "failed"]:
                logger.warning("Cannot cancel finished batch",
                              batch_id=batch_id,
                              status=batch.status)
                return batch
            
            update_data = {
                "status": "cancelled",
                "completed_at": datetime.now(timezone.utc),
                "cancelled_by": cancelled_by,
                "error_message": f"Cancelled by {cancelled_by}" if cancelled_by else "Cancelled"
            }
            
            updated_batch = await self.update(batch_id, update_data)
            
            logger.info("Batch cancelled",
                       batch_id=batch_id,
                       cancelled_by=cancelled_by)
            
            return updated_batch
            
        except Exception as e:
            logger.error("Failed to cancel batch",
                        batch_id=batch_id,
                        error=str(e))
            raise DatabaseError(f"Failed to cancel batch: {str(e)}")
    
    async def get_active_batches(self, tenant_id: str = None) -> List[PredictionBatch]:
        """Get currently active (pending/processing) batches"""
        try:
            filters = {"status": "processing"}
            if tenant_id:
                # Need to handle multiple status values with raw query
                query_text = """
                    SELECT * FROM prediction_batches
                    WHERE status IN ('pending', 'processing')
                    AND tenant_id = :tenant_id
                    ORDER BY requested_at DESC
                """
                params = {"tenant_id": tenant_id}
            else:
                query_text = """
                    SELECT * FROM prediction_batches
                    WHERE status IN ('pending', 'processing')
                    ORDER BY requested_at DESC
                """
                params = {}
            
            result = await self.session.execute(text(query_text), params)
            
            batches = []
            for row in result.fetchall():
                record_dict = dict(row._mapping)
                batch = self.model(**record_dict)
                batches.append(batch)
            
            return batches
            
        except Exception as e:
            logger.error("Failed to get active batches",
                        tenant_id=tenant_id,
                        error=str(e))
            return []
    
    async def get_batch_statistics(self, tenant_id: str = None) -> Dict[str, Any]:
        """Get batch processing statistics"""
        try:
            base_filter = "WHERE 1=1"
            params = {}
            
            if tenant_id:
                base_filter = "WHERE tenant_id = :tenant_id"
                params["tenant_id"] = tenant_id
            
            # Get counts by status
            status_query = text(f"""
                SELECT 
                    status,
                    COUNT(*) as count,
                    AVG(CASE WHEN processing_time_ms IS NOT NULL THEN processing_time_ms END) as avg_processing_time_ms
                FROM prediction_batches 
                {base_filter}
                GROUP BY status
            """)
            
            result = await self.session.execute(status_query, params)
            
            status_stats = {}
            total_batches = 0
            avg_processing_times = {}
            
            for row in result.fetchall():
                status_stats[row.status] = row.count
                total_batches += row.count
                if row.avg_processing_time_ms:
                    avg_processing_times[row.status] = float(row.avg_processing_time_ms)
            
            # Get recent activity (batches in last 7 days)
            seven_days_ago = datetime.now(timezone.utc) - timedelta(days=7)
            recent_query = text(f"""
                SELECT COUNT(*) as count 
                FROM prediction_batches 
                {base_filter}
                AND requested_at >= :seven_days_ago
            """)
            
            recent_result = await self.session.execute(recent_query, {
                **params,
                "seven_days_ago": seven_days_ago
            })
            recent_batches = recent_result.scalar() or 0
            
            # Calculate success rate
            completed = status_stats.get("completed", 0)
            failed = status_stats.get("failed", 0)
            cancelled = status_stats.get("cancelled", 0)
            finished_batches = completed + failed + cancelled
            
            success_rate = (completed / finished_batches * 100) if finished_batches > 0 else 0
            
            return {
                "total_batches": total_batches,
                "batches_by_status": status_stats,
                "success_rate": round(success_rate, 2),
                "recent_batches_7d": recent_batches,
                "avg_processing_times_ms": avg_processing_times
            }
            
        except Exception as e:
            logger.error("Failed to get batch statistics",
                        tenant_id=tenant_id,
                        error=str(e))
            return {
                "total_batches": 0,
                "batches_by_status": {},
                "success_rate": 0.0,
                "recent_batches_7d": 0,
                "avg_processing_times_ms": {}
            }
    
    async def cleanup_old_batches(self, days_old: int = 30) -> int:
        """Clean up old completed/failed batches"""
        try:
            cutoff_date = datetime.now(timezone.utc) - timedelta(days=days_old)
            
            query_text = """
                DELETE FROM prediction_batches 
                WHERE status IN ('completed', 'failed', 'cancelled')
                AND completed_at < :cutoff_date
            """
            
            result = await self.session.execute(text(query_text), {"cutoff_date": cutoff_date})
            deleted_count = result.rowcount
            
            logger.info("Cleaned up old prediction batches",
                       deleted_count=deleted_count,
                       days_old=days_old)
            
            return deleted_count
            
        except Exception as e:
            logger.error("Failed to cleanup old batches",
                        error=str(e))
            raise DatabaseError(f"Batch cleanup failed: {str(e)}")
    
    async def get_batch_details(self, batch_id: str) -> Dict[str, Any]:
        """Get detailed batch information"""
        try:
            batch = await self.get_by_id(batch_id)
            if not batch:
                return {"error": "Batch not found"}
            
            # Calculate completion percentage
            completion_percentage = 0
            if batch.total_products > 0:
                completion_percentage = (batch.completed_products / batch.total_products) * 100
            
            # Calculate elapsed time
            elapsed_time_ms = 0
            if batch.completed_at:
                elapsed_time_ms = int((batch.completed_at - batch.requested_at).total_seconds() * 1000)
            elif batch.status in ["pending", "processing"]:
                elapsed_time_ms = int((datetime.now(timezone.utc) - batch.requested_at).total_seconds() * 1000)
            
            return {
                "batch_id": str(batch.id),
                "tenant_id": str(batch.tenant_id),
                "batch_name": batch.batch_name,
                "status": batch.status,
                "progress": {
                    "total_products": batch.total_products,
                    "completed_products": batch.completed_products,
                    "failed_products": batch.failed_products,
                    "completion_percentage": round(completion_percentage, 2)
                },
                "timing": {
                    "requested_at": batch.requested_at.isoformat(),
                    "completed_at": batch.completed_at.isoformat() if batch.completed_at else None,
                    "elapsed_time_ms": elapsed_time_ms,
                    "processing_time_ms": batch.processing_time_ms
                },
                "configuration": {
                    "forecast_days": batch.forecast_days,
                    "business_type": batch.business_type
                },
                "error_message": batch.error_message,
                "cancelled_by": batch.cancelled_by
            }
            
        except Exception as e:
            logger.error("Failed to get batch details",
                        batch_id=batch_id,
                        error=str(e))
            return {"error": f"Failed to get batch details: {str(e)}"}