#!/usr/bin/env python3
"""
Portainer Integration Module
Provides enhanced Docker management through Portainer API
"""

import requests
import json
from typing import Dict, List, Optional
import os

class PortainerClient:
    def __init__(self, base_url: str, username: str = None, password: str = None):
        self.base_url = base_url.rstrip('/')
        self.username = username or os.getenv('PORTAINER_USERNAME')
        self.password = password or os.getenv('PORTAINER_PASSWORD')
        self.token = None
        self.endpoint_id = None
        
    def authenticate(self) -> bool:
        """Authenticate with Portainer API"""
        try:
            auth_url = f"{self.base_url}/api/auth"
            payload = {
                "username": self.username,
                "password": self.password
            }
            
            response = requests.post(auth_url, json=payload)
            if response.status_code == 200:
                self.token = response.json().get('jwt')
                return True
            return False
        except Exception as e:
            print(f"Portainer authentication failed: {e}")
            return False
    
    def get_endpoints(self) -> List[Dict]:
        """Get available endpoints"""
        if not self.token:
            return []
        
        try:
            headers = {'Authorization': f'Bearer {self.token}'}
            response = requests.get(f"{self.base_url}/api/endpoints", headers=headers)
            
            if response.status_code == 200:
                return response.json()
            return []
        except Exception as e:
            print(f"Error fetching endpoints: {e}")
            return []
    
    def get_containers(self, endpoint_id: int = None) -> List[Dict]:
        """Get containers from Portainer"""
        if not self.token:
            return []
        
        try:
            if not endpoint_id:
                endpoints = self.get_endpoints()
                if endpoints:
                    endpoint_id = endpoints[0]['Id']
                else:
                    return []
            
            headers = {'Authorization': f'Bearer {self.token}'}
            url = f"{self.base_url}/api/endpoints/{endpoint_id}/docker/containers/json?all=true"
            
            response = requests.get(url, headers=headers)
            
            if response.status_code == 200:
                containers = response.json()
                return self._format_portainer_containers(containers)
            return []
        except Exception as e:
            print(f"Error fetching containers from Portainer: {e}")
            return []
    
    def _format_portainer_containers(self, containers: List[Dict]) -> List[Dict]:
        """Format Portainer containers to match our standard format"""
        formatted = []
        
        for container in containers:
            formatted_container = {
                'id': container.get('Id', '')[:12],
                'name': container.get('Names', [''])[0].lstrip('/'),
                'status': container.get('State', 'unknown'),
                'image': container.get('Image', ''),
                'ports': self._format_ports(container.get('Ports', [])),
                'created': container.get('Created', ''),
                'uptime': self._calculate_uptime(container),
                'labels': container.get('Labels', {}),
                'environment': [],
                'mounts': self._format_mounts(container.get('Mounts', [])),
                'networks': self._format_networks(container.get('NetworkSettings', {}).get('Networks', {}))
            }
            formatted.append(formatted_container)
        
        return formatted
    
    def _format_ports(self, ports: List[Dict]) -> List[Dict]:
        """Format port mappings"""
        formatted_ports = []
        for port in ports:
            formatted_ports.append({
                'container_port': f"{port.get('PrivatePort', '')}/{port.get('Type', 'tcp')}",
                'host_ip': port.get('IP', '0.0.0.0'),
                'host_port': str(port.get('PublicPort', ''))
            })
        return formatted_ports
    
    def _format_mounts(self, mounts: List[Dict]) -> List[Dict]:
        """Format volume mounts"""
        formatted_mounts = []
        for mount in mounts:
            formatted_mounts.append({
                'type': mount.get('Type'),
                'source': mount.get('Source'),
                'destination': mount.get('Destination'),
                'mode': mount.get('Mode')
            })
        return formatted_mounts
    
    def _format_networks(self, networks: Dict) -> List[Dict]:
        """Format network information"""
        formatted_networks = []
        for network_name, network_config in networks.items():
            formatted_networks.append({
                'name': network_name,
                'ip_address': network_config.get('IPAddress'),
                'gateway': network_config.get('Gateway'),
                'mac_address': network_config.get('MacAddress')
            })
        return formatted_networks
    
    def _calculate_uptime(self, container: Dict) -> str:
        """Calculate uptime from container info"""
        status = container.get('State')
        if status != 'running':
            return 'Down'
        
        # For now, return simplified uptime
        # In a real implementation, you'd parse the Created field
        return 'Running'
    
    def get_stacks(self, endpoint_id: int = None) -> List[Dict]:
        """Get Docker stacks from Portainer"""
        if not self.token:
            return []
        
        try:
            if not endpoint_id:
                endpoints = self.get_endpoints()
                if endpoints:
                    endpoint_id = endpoints[0]['Id']
                else:
                    return []
            
            headers = {'Authorization': f'Bearer {self.token}'}
            url = f"{self.base_url}/api/endpoints/{endpoint_id}/docker/stacks"
            
            response = requests.get(url, headers=headers)
            
            if response.status_code == 200:
                return response.json()
            return []
        except Exception as e:
            print(f"Error fetching stacks from Portainer: {e}")
            return []
    
    def get_volumes(self, endpoint_id: int = None) -> List[Dict]:
        """Get Docker volumes from Portainer"""
        if not self.token:
            return []
        
        try:
            if not endpoint_id:
                endpoints = self.get_endpoints()
                if endpoints:
                    endpoint_id = endpoints[0]['Id']
                else:
                    return []
            
            headers = {'Authorization': f'Bearer {self.token}'}
            url = f"{self.base_url}/api/endpoints/{endpoint_id}/docker/volumes"
            
            response = requests.get(url, headers=headers)
            
            if response.status_code == 200:
                volumes_data = response.json()
                return volumes_data.get('Volumes', [])
            return []
        except Exception as e:
            print(f"Error fetching volumes from Portainer: {e}")
            return []
    
    def get_networks(self, endpoint_id: int = None) -> List[Dict]:
        """Get Docker networks from Portainer"""
        if not self.token:
            return []
        
        try:
            if not endpoint_id:
                endpoints = self.get_endpoints()
                if endpoints:
                    endpoint_id = endpoints[0]['Id']
                else:
                    return []
            
            headers = {'Authorization': f'Bearer {self.token}'}
            url = f"{self.base_url}/api/endpoints/{endpoint_id}/docker/networks"
            
            response = requests.get(url, headers=headers)
            
            if response.status_code == 200:
                return response.json()
            return []
        except Exception as e:
            print(f"Error fetching networks from Portainer: {e}")
            return []
    
    def start_container(self, container_id: str, endpoint_id: int = None) -> bool:
        """Start a container"""
        return self._container_action(container_id, 'start', endpoint_id)
    
    def stop_container(self, container_id: str, endpoint_id: int = None) -> bool:
        """Stop a container"""
        return self._container_action(container_id, 'stop', endpoint_id)
    
    def restart_container(self, container_id: str, endpoint_id: int = None) -> bool:
        """Restart a container"""
        return self._container_action(container_id, 'restart', endpoint_id)
    
    def _container_action(self, container_id: str, action: str, endpoint_id: int = None) -> bool:
        """Perform container action"""
        if not self.token:
            return False
        
        try:
            if not endpoint_id:
                endpoints = self.get_endpoints()
                if endpoints:
                    endpoint_id = endpoints[0]['Id']
                else:
                    return False
            
            headers = {'Authorization': f'Bearer {self.token}'}
            url = f"{self.base_url}/api/endpoints/{endpoint_id}/docker/containers/{container_id}/{action}"
            
            response = requests.post(url, headers=headers)
            return response.status_code == 204
        except Exception as e:
            print(f"Error performing {action} on container: {e}")
            return False