"""
MongoDB database setup and models for storing classification results.
Uses Motor (async MongoDB driver) for FastAPI integration.
"""
import time
import os
import base64
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any
from motor.motor_asyncio import AsyncIOMotorClient
from pymongo.errors import ConnectionFailure
import logging

from .config import settings

logger = logging.getLogger(__name__)

# MongoDB connection settings
MONGODB_URL = settings.MONGODB_URL
DATABASE_NAME = settings.MONGODB_DATABASE

# Global MongoDB client
mongodb_client: Optional[AsyncIOMotorClient] = None
mongodb_db = None
gridfs_bucket = None


async def connect_to_mongodb():
    """Connect to MongoDB database."""
    global mongodb_client, mongodb_db, gridfs_bucket
    
    if not MONGODB_URL:
        raise ValueError(
            "MONGODB_URL environment variable is not set. "
            "Please set it to your MongoDB connection string."
        )
    
    try:
        from motor.motor_asyncio import AsyncIOMotorGridFSBucket
        
        mongodb_client = AsyncIOMotorClient(MONGODB_URL)
        # Test connection
        await mongodb_client.admin.command('ping')
        mongodb_db = mongodb_client[DATABASE_NAME]
        
        # Initialize GridFS bucket for large file storage (videos)
        gridfs_bucket = AsyncIOMotorGridFSBucket(mongodb_db)
        
        logger.info(f"Connected to MongoDB database: {DATABASE_NAME}")
        return mongodb_db
    except ConnectionFailure as e:
        logger.error(f"Failed to connect to MongoDB: {e}")
        raise


async def close_mongodb_connection():
    """Close MongoDB connection."""
    global mongodb_client
    if mongodb_client:
        mongodb_client.close()
        logger.info("MongoDB connection closed")


def get_database():
    """Get MongoDB database instance."""
    if mongodb_db is None:
        raise RuntimeError("MongoDB database not initialized. Call connect_to_mongodb() first.")
    return mongodb_db


def get_gridfs_bucket():
    """Get GridFS bucket for large file storage."""
    if gridfs_bucket is None:
        raise RuntimeError("GridFS bucket not initialized. Call connect_to_mongodb() first.")
    return gridfs_bucket


async def create_indexes():
    """Create database indexes for better query performance."""
    db = get_database()
    
    # Image classifications collection
    collection = db["classifications"]
    await collection.create_index("id", unique=True)
    await collection.create_index("timestamp")
    await collection.create_index([("latitude", 1), ("longitude", 1)])
    
    # Video classifications collection
    video_collection = db["video_classifications"]
    await video_collection.create_index("id", unique=True)
    await video_collection.create_index("timestamp")
    await video_collection.create_index([("latitude", 1), ("longitude", 1)])
    
    # Job tracking collection
    jobs_collection = db["jobs"]
    await jobs_collection.create_index("job_id", unique=True)
    await jobs_collection.create_index("status")
    await jobs_collection.create_index("created_at")
    
    logger.info("MongoDB indexes created successfully")


class ClassificationRecord:
    """
    MongoDB document model for classification results.
    """
    
    @staticmethod
    def encode_image(image_bytes: bytes) -> str:
        """Encode image bytes to base64 string for MongoDB storage."""
        return base64.b64encode(image_bytes).decode('utf-8')
    
    @staticmethod
    def decode_image(image_base64: str) -> bytes:
        """Decode base64 string back to image bytes."""
        return base64.b64decode(image_base64.encode('utf-8'))
    
    @staticmethod
    async def create(
        record_id: str,
        image_bytes: bytes,
        latitude: Optional[float],
        longitude: Optional[float],
        flowers: List[Dict[str, Any]],
        filename: str = "image.jpg"
    ) -> Dict[str, Any]:
        """Create a new classification record in MongoDB.
        
        Note: Only stores the annotated image (with bounding boxes), not the original.
        """
        db = get_database()
        collection = db["classifications"]
        
        # Encode annotated image to base64
        image_base64 = ClassificationRecord.encode_image(image_bytes)
        
        document = {
            "id": record_id,
            "image_base64": image_base64,
            "image_filename": filename,
            "image_content_type": "image/jpeg",
            "latitude": latitude,
            "longitude": longitude,
            "timestamp": datetime.utcnow(),
            "flowers": flowers,
        }
        
        result = await collection.insert_one(document)
        return document
    
    @staticmethod
    async def get_by_id(record_id: str) -> Optional[Dict[str, Any]]:
        """Get classification record by ID."""
        db = get_database()
        collection = db["classifications"]
        return await collection.find_one({"id": record_id})

    @staticmethod
    async def get_all() -> List[Dict[str, Any]]:
        db = get_database()
        collection = db["classifications"]

        cursor = (
            collection
            .find(
                {},
                {
                    "id": 1,
                    "flowers": 1,
                    "timestamp": 1,
                    "_id": 0,   # exclude Mongo ObjectId
                },
            )
            .sort("timestamp", -1)
        )

        return await cursor.to_list(length=None) 
    
    
    @staticmethod
    async def delete_by_id(record_id: str) -> bool:
        """Delete classification record by ID."""
        db = get_database()
        collection = db["classifications"]
        result = await collection.delete_one({"id": record_id})
        return result.deleted_count > 0
    
    @staticmethod
    async def delete_all() -> int:
        """Delete all classification records."""
        db = get_database()
        collection = db["classifications"]
        result = await collection.delete_many({})
        return result.deleted_count

    @staticmethod
    async def count() -> int:
        """Get total number of classification records."""
        db = get_database()
        collection = db["classifications"]
        return await collection.count_documents({})


class VideoClassificationRecord:
    """
    MongoDB document model for video classification results.
    Uses GridFS for storing large video files.
    """
    
    @staticmethod
    async def create(
        record_id: str,
        video_bytes: bytes,
        latitude: Optional[float],
        longitude: Optional[float],
        frame_results: List[List[Dict[str, Any]]],
        video_metadata: Dict[str, Any],
        filename: str = "video.mp4"
    ) -> Dict[str, Any]:
        """
        Create a new video classification record in MongoDB.
        
        Stores video in GridFS (for large files) and metadata in regular collection.
        
        Args:
            record_id: Unique identifier for the record
            video_bytes: Raw video bytes (annotated)
            latitude: GPS latitude
            longitude: GPS longitude
            frame_results: List of detection lists (one per frame)
            video_metadata: Dict with 'total_frames', 'fps', 'duration'
            filename: Original filename
            
        Returns:
            Created document
        """
        db = get_database()
        bucket = get_gridfs_bucket()
        collection = db["video_classifications"]
        
        # Store video in GridFS
        video_id = await bucket.upload_from_stream(
            filename,
            video_bytes,
            metadata={
                "record_id": record_id,
                "content_type": "video/mp4",
                "uploaded_at": datetime.utcnow()
            }
        )
        
        # Aggregate statistics across all frames
        stage_summary = {0: 0, 1: 0, 2: 0}
        total_detections = 0
        
        for frame_detections in frame_results:
            total_detections += len(frame_detections)
            for det in frame_detections:
                stage = det['stage']
                stage_summary[stage] = stage_summary.get(stage, 0) + 1
        
        # Convert integer keys to strings for MongoDB compatibility
        stage_summary = {str(k): v for k, v in stage_summary.items()}
        
        # Calculate average flowers per frame
        total_frames = video_metadata.get('total_frames', len(frame_results))
        avg_flowers = total_detections / total_frames if total_frames > 0 else 0
        
        # Store metadata in regular collection
        document = {
            "id": record_id,
            "video_gridfs_id": video_id,
            "video_filename": filename,
            "video_content_type": "video/mp4",
            "latitude": latitude,
            "longitude": longitude,
            "timestamp": datetime.utcnow(),
            "total_frames": total_frames,
            "fps": video_metadata.get('fps', 0),
            "duration_seconds": video_metadata.get('duration', 0),
            "frame_results": frame_results,  # Store all frame-by-frame detections
            "stage_summary": stage_summary,
            "total_detections": total_detections,
            "average_flowers_per_frame": avg_flowers,
        }
        
        await collection.insert_one(document)
        logger.info(f"Video classification record created: {record_id}")
        return document
    
    @staticmethod
    async def get_by_id(record_id: str) -> Optional[Dict[str, Any]]:
        """Get video classification record by ID."""
        db = get_database()
        collection = db["video_classifications"]
        return await collection.find_one({"id": record_id})
    
    @staticmethod
    async def get_all() -> List[Dict[str, Any]]:
        """Get all video classification records, sorted by timestamp."""
        db = get_database()
        collection = db["video_classifications"]
        cursor = collection.find({}).sort("timestamp", -1)
        return await cursor.to_list(length=None)
    
    @staticmethod
    async def get_video_bytes(record_id: str) -> Optional[bytes]:
        """
        Retrieve video bytes from GridFS.
        
        Args:
            record_id: Record identifier
            
        Returns:
            Video bytes or None if not found
        """
        record = await VideoClassificationRecord.get_by_id(record_id)
        if not record:
            return None
        
        bucket = get_gridfs_bucket()
        video_id = record.get('video_gridfs_id')
        
        if not video_id:
            logger.error(f"No video_gridfs_id found for record {record_id}")
            return None
        
        try:
            # Download from GridFS
            grid_out = await bucket.open_download_stream(video_id)
            video_bytes = await grid_out.read()
            return video_bytes
        except Exception as e:
            logger.error(f"Error downloading video from GridFS: {e}")
            return None
    
    @staticmethod
    async def delete_by_id(record_id: str) -> bool:
        """
        Delete video classification record by ID.
        Also deletes the video file from GridFS.
        """
        db = get_database()
        bucket = get_gridfs_bucket()
        collection = db["video_classifications"]
        
        # Get record first to get GridFS ID
        record = await VideoClassificationRecord.get_by_id(record_id)
        if not record:
            return False
        
        # Delete video from GridFS
        video_id = record.get('video_gridfs_id')
        if video_id:
            try:
                await bucket.delete(video_id)
                logger.info(f"Deleted video from GridFS: {video_id}")
            except Exception as e:
                logger.error(f"Error deleting video from GridFS: {e}")
        
        # Delete metadata document
        result = await collection.delete_one({"id": record_id})
        return result.deleted_count > 0
    
    @staticmethod
    async def delete_all() -> int:
        """Delete all video classification records and their GridFS video files."""
        db = get_database()
        bucket = get_gridfs_bucket()
        collection = db["video_classifications"]

        records = await collection.find({}, {"video_gridfs_id": 1}).to_list(length=None)
        for record in records:
            video_id = record.get("video_gridfs_id")
            if video_id:
                try:
                    await bucket.delete(video_id)
                except Exception as e:
                    logger.error(f"Error deleting video from GridFS: {e}")

        result = await collection.delete_many({})
        return result.deleted_count

    @staticmethod
    async def count() -> int:
        """Get total number of video classification records."""
        db = get_database()
        collection = db["video_classifications"]
        return await collection.count_documents({})


class JobRecord:
    """
    MongoDB document model for async job tracking.
    """
    
    @staticmethod
    async def create(
        job_id: str,
        job_type: str,
        status: str = "queued",
        metadata: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Create a new job record."""
        db = get_database()
        collection = db["jobs"]
        
        document = {
            "job_id": job_id,
            "job_type": job_type,
            "status": status,
            "progress": 0,
            "message": None,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow(),
            "completed_at": None,
            "metadata": metadata or {},
            "result": None,
            "error": None,
        }
        
        await collection.insert_one(document)
        logger.info(f"Job record created: {job_id}")
        return document
    
    @staticmethod
    async def update_status(
        job_id: str,
        status: str,
        progress: Optional[int] = None,
        message: Optional[str] = None,
        result: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None
    ) -> bool:
        """Update job status."""
        db = get_database()
        collection = db["jobs"]
        
        update_data = {
            "status": status,
            "updated_at": datetime.utcnow(),
        }
        
        if progress is not None:
            update_data["progress"] = progress
        if message is not None:
            update_data["message"] = message
        if result is not None:
            update_data["result"] = result
        if error is not None:
            update_data["error"] = error
        
        if status in ["completed", "failed"]:
            update_data["completed_at"] = datetime.utcnow()
        
        result = await collection.update_one(
            {"job_id": job_id},
            {"$set": update_data}
        )
        
        return result.modified_count > 0
    
    @staticmethod
    async def get_by_id(job_id: str) -> Optional[Dict[str, Any]]:
        """Get job record by ID."""
        db = get_database()
        collection = db["jobs"]
        return await collection.find_one({"job_id": job_id})
    
    @staticmethod
    async def get_active_jobs() -> List[Dict[str, Any]]:
        """Get all active (not completed/failed) jobs."""
        db = get_database()
        collection = db["jobs"]
        cursor = collection.find({
            "status": {"$in": ["queued", "processing"]}
        }).sort("created_at", -1)
        return await cursor.to_list(length=None)
    
    @staticmethod
    async def delete_by_id(job_id: str) -> bool:
        """Delete job record by ID."""
        db = get_database()
        collection = db["jobs"]
        result = await collection.delete_one({"job_id": job_id})
        return result.deleted_count > 0
    
    @staticmethod
    async def delete_all() -> int:
        """Delete all job records."""
        db = get_database()
        collection = db["jobs"]
        result = await collection.delete_many({})
        return result.deleted_count

    @staticmethod
    async def cleanup_old_jobs(days: int = 7) -> int:
        """Delete job records older than specified days."""
        db = get_database()
        collection = db["jobs"]
        cutoff_date = datetime.utcnow() - timedelta(days=days)
        result = await collection.delete_many({
            "created_at": {"$lt": cutoff_date},
            "status": {"$in": ["completed", "failed"]}
        })
        logger.info(f"Cleaned up {result.deleted_count} old job records")
        return result.deleted_count