#!/bin/bash

# PyGuardian Agent Installation Script
# Usage: ./install_agent.sh --master <master_ip> --port <master_port>

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Default values
MASTER_IP=""
MASTER_PORT="8080"
AGENT_PORT="8081"
INSTALL_DIR="/opt/pyguardian-agent"
SERVICE_NAME="pyguardian-agent"

# Functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

show_usage() {
    echo "PyGuardian Agent Installation Script"
    echo ""
    echo "Usage: $0 --master <master_ip> [OPTIONS]"
    echo ""
    echo "Required:"
    echo "  --master <ip>      Master server IP address"
    echo ""
    echo "Optional:"
    echo "  --port <port>      Master server port (default: 8080)"
    echo "  --agent-port <p>   Agent listen port (default: 8081)"
    echo "  --help             Show this help"
    echo ""
    echo "Example:"
    echo "  $0 --master 192.168.1.100 --port 8080"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --master)
            MASTER_IP="$2"
            shift 2
            ;;
        --port)
            MASTER_PORT="$2"
            shift 2
            ;;
        --agent-port)
            AGENT_PORT="$2"
            shift 2
            ;;
        --help)
            show_usage
            exit 0
            ;;
        *)
            log_error "Unknown option: $1"
            show_usage
            exit 1
            ;;
    esac
done

# Validate required parameters
if [[ -z "$MASTER_IP" ]]; then
    log_error "Master IP is required"
    show_usage
    exit 1
fi

# Check if running as root
if [[ $EUID -ne 0 ]]; then
    log_error "This script must be run as root"
    exit 1
fi

log_info "Starting PyGuardian Agent installation..."
log_info "Master Server: ${MASTER_IP}:${MASTER_PORT}"
log_info "Agent Port: ${AGENT_PORT}"

# Check system requirements
log_info "Checking system requirements..."

# Check Python version
if ! command -v python3 &> /dev/null; then
    log_error "Python3 is required but not installed"
    exit 1
fi

PYTHON_VERSION=$(python3 -c 'import sys; print(".".join(map(str, sys.version_info[:2])))')
REQUIRED_VERSION="3.10"

if ! python3 -c "import sys; sys.exit(0 if sys.version_info >= (3, 10) else 1)" 2>/dev/null; then
    log_error "Python 3.10+ is required. Found: $PYTHON_VERSION"
    exit 1
fi

log_success "Python version check passed: $PYTHON_VERSION"

# Install system dependencies
log_info "Installing system dependencies..."

if command -v apt-get &> /dev/null; then
    apt-get update
    apt-get install -y python3-pip python3-venv curl wget
elif command -v yum &> /dev/null; then
    yum install -y python3-pip curl wget
elif command -v dnf &> /dev/null; then
    dnf install -y python3-pip curl wget
else
    log_warning "Unknown package manager. Please install python3-pip manually."
fi

# Create installation directory
log_info "Creating installation directory..."
mkdir -p "$INSTALL_DIR"
mkdir -p /etc/pyguardian-agent
mkdir -p /var/log/pyguardian-agent

# Create agent configuration
log_info "Creating agent configuration..."
cat > /etc/pyguardian-agent/config.yaml << EOF
# PyGuardian Agent Configuration

agent:
  port: ${AGENT_PORT}
  master_host: "${MASTER_IP}"
  master_port: ${MASTER_PORT}
  heartbeat_interval: 30
  max_reconnect_attempts: 10
  reconnect_delay: 5

logging:
  level: INFO
  file: /var/log/pyguardian-agent/agent.log
  max_size: 10MB
  backup_count: 5

security:
  ssl_verify: true
  api_key_file: /etc/pyguardian-agent/api.key
EOF

# Create simple agent script
log_info "Creating agent script..."
cat > "$INSTALL_DIR/agent.py" << 'EOF'
#!/usr/bin/env python3
"""
PyGuardian Agent
Lightweight agent for cluster management
"""

import asyncio
import json
import logging
import signal
import sys
import time
import yaml
import aiohttp
import psutil
from pathlib import Path

class PyGuardianAgent:
    def __init__(self, config_path="/etc/pyguardian-agent/config.yaml"):
        self.config_path = config_path
        self.config = self.load_config()
        self.setup_logging()
        self.session = None
        self.running = False
        
    def load_config(self):
        """Load agent configuration"""
        try:
            with open(self.config_path, 'r') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"Error loading config: {e}")
            sys.exit(1)
    
    def setup_logging(self):
        """Setup logging configuration"""
        log_file = self.config.get('logging', {}).get('file', '/var/log/pyguardian-agent/agent.log')
        log_level = getattr(logging, self.config.get('logging', {}).get('level', 'INFO'))
        
        # Create log directory if it doesn't exist
        Path(log_file).parent.mkdir(parents=True, exist_ok=True)
        
        logging.basicConfig(
            level=log_level,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)
    
    async def get_system_info(self):
        """Get system information"""
        try:
            cpu_percent = psutil.cpu_percent(interval=1)
            memory = psutil.virtual_memory()
            disk = psutil.disk_usage('/')
            
            return {
                'status': 'online',
                'cpu_percent': cpu_percent,
                'memory_percent': memory.percent,
                'memory_total': memory.total,
                'memory_used': memory.used,
                'disk_percent': disk.percent,
                'disk_total': disk.total,
                'disk_used': disk.used,
                'timestamp': time.time()
            }
        except Exception as e:
            self.logger.error(f"Error getting system info: {e}")
            return {'status': 'error', 'error': str(e)}
    
    async def send_heartbeat(self):
        """Send heartbeat to master server"""
        try:
            system_info = await self.get_system_info()
            
            master_url = f"http://{self.config['agent']['master_host']}:{self.config['agent']['master_port']}"
            
            async with self.session.post(
                f"{master_url}/api/agent/heartbeat",
                json=system_info,
                timeout=10
            ) as response:
                if response.status == 200:
                    self.logger.debug("Heartbeat sent successfully")
                else:
                    self.logger.warning(f"Heartbeat failed with status: {response.status}")
                    
        except Exception as e:
            self.logger.error(f"Error sending heartbeat: {e}")
    
    async def heartbeat_loop(self):
        """Main heartbeat loop"""
        interval = self.config.get('agent', {}).get('heartbeat_interval', 30)
        
        while self.running:
            await self.send_heartbeat()
            await asyncio.sleep(interval)
    
    async def start(self):
        """Start the agent"""
        self.logger.info("Starting PyGuardian Agent...")
        self.running = True
        
        # Create HTTP session
        self.session = aiohttp.ClientSession()
        
        try:
            # Start heartbeat loop
            await self.heartbeat_loop()
        except Exception as e:
            self.logger.error(f"Agent error: {e}")
        finally:
            await self.stop()
    
    async def stop(self):
        """Stop the agent"""
        self.logger.info("Stopping PyGuardian Agent...")
        self.running = False
        
        if self.session:
            await self.session.close()

async def main():
    agent = PyGuardianAgent()
    
    # Handle signals
    def signal_handler(signum, frame):
        print(f"\nReceived signal {signum}, shutting down...")
        asyncio.create_task(agent.stop())
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    try:
        await agent.start()
    except KeyboardInterrupt:
        await agent.stop()

if __name__ == "__main__":
    asyncio.run(main())
EOF

chmod +x "$INSTALL_DIR/agent.py"

# Install Python dependencies
log_info "Installing Python dependencies..."
python3 -m pip install --upgrade pip
python3 -m pip install aiohttp pyyaml psutil

# Create systemd service
log_info "Creating systemd service..."
cat > "/etc/systemd/system/${SERVICE_NAME}.service" << EOF
[Unit]
Description=PyGuardian Agent
After=network.target
Wants=network.target

[Service]
Type=simple
User=root
Group=root
WorkingDirectory=${INSTALL_DIR}
ExecStart=/usr/bin/python3 ${INSTALL_DIR}/agent.py
ExecReload=/bin/kill -HUP \$MAINPID
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF

# Reload systemd and enable service
log_info "Enabling systemd service..."
systemctl daemon-reload
systemctl enable "$SERVICE_NAME"

# Start the service
log_info "Starting PyGuardian Agent..."
systemctl start "$SERVICE_NAME"

# Check service status
sleep 3
if systemctl is-active --quiet "$SERVICE_NAME"; then
    log_success "PyGuardian Agent installed and started successfully!"
    log_info "Service status: $(systemctl is-active $SERVICE_NAME)"
    log_info "Check logs: journalctl -u $SERVICE_NAME -f"
    log_info "Agent config: /etc/pyguardian-agent/config.yaml"
    log_info "Agent logs: /var/log/pyguardian-agent/agent.log"
else
    log_error "PyGuardian Agent failed to start"
    log_info "Check service status: systemctl status $SERVICE_NAME"
    log_info "Check logs: journalctl -u $SERVICE_NAME"
    exit 1
fi

log_success "Installation completed!"
log_info "The agent should now be visible in your PyGuardian master server."
log_info "Use '/agents' command in Telegram to verify the agent connection."
EOF