from typing import Dict, List, Optional
from datetime import datetime
import uuid
import logging

logger = logging.getLogger(__name__)

class FirmwareManager:
    """Firmware version tracking and upgrade management"""
    
    def __init__(self, db):
        self.db = db
    
    async def register_firmware(self, vendor: str, model: str, version: str,
                               file_url: Optional[str] = None,
                               release_date: Optional[datetime] = None,
                               notes: Optional[str] = None,
                               org_id: str = None) -> Dict:
        """Register firmware version"""
        firmware = {
            'id': str(uuid.uuid4()),
            'vendor': vendor,
            'model': model,
            'version': version,
            'file_url': file_url,
            'release_date': release_date or datetime.utcnow(),
            'registered_at': datetime.utcnow(),
            'notes': notes,
            'status': 'available',
        }
        if org_id:
            firmware['organization_id'] = org_id
        
        await self.db.firmwares.insert_one(firmware)
        
        logger.info(f"Registered firmware: {vendor} {model} {version}")
        return firmware
    
    async def get_firmware_list(self, vendor: Optional[str] = None,
                               model: Optional[str] = None,
                               org_id: str = None) -> List[Dict]:
        """Get firmware list"""
        query = {}
        if org_id:
            query['organization_id'] = org_id
        if vendor:
            query['vendor'] = vendor
        if model:
            query['model'] = model
        
        firmwares = await self.db.firmwares.find(query).sort('release_date', -1).to_list(100)
        return firmwares
    
    async def get_latest_firmware(self, vendor: str, model: str) -> Optional[Dict]:
        """Get latest firmware for device"""
        firmware = await self.db.firmwares.find_one(
            {'vendor': vendor, 'model': model},
            sort=[('release_date', -1)]
        )
        return firmware
    
    async def check_firmware_updates(self, device_id: str, org_id: str = None) -> Dict:
        """Check if firmware update is available"""
        query = {'id': device_id}
        if org_id:
            query['organization_id'] = org_id
        device = await self.db.devices.find_one(query)
        if not device:
            raise Exception("Device not found")
        
        vendor = device.get('vendor')
        model = device.get('model')
        current_version = device.get('firmware_version', 'Unknown')
        
        if not vendor or not model:
            return {
                'update_available': False,
                'message': 'Device vendor/model information not available'
            }
        
        latest_firmware = await self.get_latest_firmware(vendor, model)
        
        if not latest_firmware:
            return {
                'update_available': False,
                'current_version': current_version,
                'message': 'No firmware found for this device'
            }
        
        update_available = latest_firmware['version'] != current_version
        
        return {
            'update_available': update_available,
            'current_version': current_version,
            'latest_version': latest_firmware['version'],
            'release_date': latest_firmware['release_date'],
            'notes': latest_firmware.get('notes', ''),
            'firmware_id': latest_firmware['id']
        }
    
    async def schedule_firmware_upgrade(self, device_id: str, firmware_id: str,
                                       scheduled_time: Optional[datetime] = None,
                                       auto_reboot: bool = True,
                                       org_id: str = None) -> Dict:
        """Schedule firmware upgrade scoped to organization"""
        query = {'id': device_id}
        if org_id:
            query['organization_id'] = org_id
        device = await self.db.devices.find_one(query)
        if not device:
            raise Exception("Device not found")
        
        firmware = await self.db.firmwares.find_one({'id': firmware_id})
        if not firmware:
            raise Exception("Firmware not found")
        
        upgrade_job = {
            'id': str(uuid.uuid4()),
            'device_id': device_id,
            'device_ip': device['ip'],
            'firmware_id': firmware_id,
            'firmware_version': firmware['version'],
            'scheduled_time': scheduled_time or datetime.utcnow(),
            'auto_reboot': auto_reboot,
            'status': 'scheduled',
            'created_at': datetime.utcnow(),
            'started_at': None,
            'completed_at': None,
            'error_message': None
        }
        if org_id:
            upgrade_job['organization_id'] = org_id
        
        await self.db.firmware_upgrades.insert_one(upgrade_job)
        
        logger.info(f"Scheduled firmware upgrade for device {device_id}")
        return upgrade_job
    
    async def get_upgrade_history(self, device_id: str, org_id: str = None) -> List[Dict]:
        """Get firmware upgrade history scoped to organization"""
        query = {'device_id': device_id}
        if org_id:
            query['organization_id'] = org_id
        upgrades = await self.db.firmware_upgrades.find(query).sort('created_at', -1).to_list(50)
        
        return upgrades
    
    async def update_device_firmware_info(self, device_id: str, 
                                         firmware_version: str,
                                         firmware_date: Optional[datetime] = None,
                                         org_id: str = None):
        """Update device firmware information scoped to organization"""
        update_data = {
            'firmware_version': firmware_version,
            'firmware_updated_at': datetime.utcnow()
        }
        
        if firmware_date:
            update_data['firmware_release_date'] = firmware_date
        
        query = {'id': device_id}
        if org_id:
            query['organization_id'] = org_id
        
        await self.db.devices.update_one(query, {'$set': update_data})
        
        logger.info(f"Updated firmware info for device {device_id}: {firmware_version}")
    
    async def get_devices_with_outdated_firmware(self, org_id: str = None) -> List[Dict]:
        """Get devices that have outdated firmware scoped to organization"""
        query = {
            'vendor': {'$exists': True},
            'model': {'$exists': True}
        }
        if org_id:
            query['organization_id'] = org_id
        
        devices = await self.db.devices.find(query).to_list(1000)
        
        outdated_devices = []
        
        for device in devices:
            check_result = await self.check_firmware_updates(device['id'], org_id=org_id)
            if check_result.get('update_available'):
                outdated_devices.append({
                    'device_id': device['id'],
                    'ip': device['ip'],
                    'hostname': device.get('hostname', device['ip']),
                    'vendor': device.get('vendor'),
                    'model': device.get('model'),
                    'current_version': check_result['current_version'],
                    'latest_version': check_result['latest_version']
                })
        
        return outdated_devices