#!/usr/bin/env python3

import argparse
import asyncio
import json
import os
import signal
import sys
import tempfile
import time
from pathlib import Path
from typing import Dict, List, Optional

# ANSI color codes for terminal output if needed
class Colors:
    COLORS = [
        '\033[94m',   # BLUE
        '\033[92m',   # GREEN
        '\033[95m',   # MAGENTA
        '\033[93m',   # YELLOW
        '\033[96m',   # CYAN
        '\033[91m',   # RED
        '\033[38;5;208m',  # ORANGE
        '\033[38;5;165m',  # PURPLE
        '\033[38;5;39m',   # DEEP BLUE
        '\033[38;5;82m',   # LIME
    ]
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    
    # Additional colors for output
    GREY = '\033[37m'
    WHITE = '\033[97m'
    RED = '\033[91m'
    ORANGE = '\033[38;5;208m'  # Using 256-color code for orange
    CYAN = '\033[96m'
    MAGENTA = '\033[95m'

class BotManager:
    """Manages multiple Icecast IRC bot instances."""
    
    def __init__(self):
        self.bots: Dict[str, asyncio.subprocess.Process] = {}
        self.config_dir = Path('.')  # Use current directory instead of /etc/icecast-irc-bot
        self.socket_dir = Path(tempfile.gettempdir())
        self.state_file = Path(tempfile.gettempdir()) / 'icecast-irc-bots.json'
        self.monitor_task = None
        self.venv_python = os.getenv('VIRTUAL_ENV', '/opt/icecast-irc-bot/venv') + '/bin/python3'
        # If venv_python doesn't exist, use the current Python interpreter
        if not os.path.exists(self.venv_python):
            self.venv_python = sys.executable
        
    def _save_state(self):
        """Save the current state of running bots."""
        try:
            # Load existing state
            existing_state = self._load_state()
            
            # Update with current bots
            current_state = {
                bot_id: {
                    'pid': process.pid,
                    'config': str(self.config_dir / f"{bot_id}.yaml")
                }
                for bot_id, process in self.bots.items()
            }
            
            # Merge states, with current bots taking precedence
            merged_state = {**existing_state, **current_state}
            
            # Save merged state
            with open(self.state_file, 'w') as f:
                json.dump(merged_state, f)
                
            print(f"Saved state with {len(merged_state)} bots")
        except Exception as e:
            print(f"Error saving state: {e}")
            pass
            
    def _load_state(self):
        """Load the state of running bots."""
        try:
            if self.state_file.exists():
                with open(self.state_file, 'r') as f:
                    state = json.load(f)
                return state
            return {}
        except Exception:
            return {}
        
    async def monitor_processes(self):
        """Monitor running bot processes and clean up dead ones."""
        while True:
            try:
                for bot_id in list(self.bots.keys()):
                    process = self.bots[bot_id]
                    try:
                        # Check if process exists in system
                        os.kill(process.pid, 0)
                        # Check if process has terminated
                        if process.returncode is not None:
                            await self._cleanup_process(process)
                            del self.bots[bot_id]
                    except ProcessLookupError:
                        await self._cleanup_process(process)
                        del self.bots[bot_id]
                    except Exception:
                        pass
            except Exception:
                pass
            await asyncio.sleep(5)  # Check every 5 seconds
            
    async def _cleanup_process(self, process: asyncio.subprocess.Process):
        """Clean up a process and its resources.
        
        Args:
            process: The process to clean up
        """
        try:
            print(f"Cleaning up process with PID {process.pid}")
            # Wait for the process to finish if it hasn't
            if process.returncode is None:
                try:
                    print(f"Terminating process {process.pid}")
                    process.terminate()
                    await asyncio.wait_for(process.wait(), timeout=5.0)
                except (asyncio.TimeoutError, ProcessLookupError) as e:
                    print(f"Error terminating process: {e}")
                    try:
                        print(f"Killing process {process.pid}")
                        process.kill()
                        await process.wait()
                    except ProcessLookupError as e:
                        print(f"Process already gone: {e}")
                        pass
                    
            # Drain any remaining output to prevent deadlocks
            if process.stdout:
                try:
                    data = await process.stdout.read()
                    if data:
                        print(f"Remaining stdout: {data.decode().strip()}")
                except (ValueError, IOError) as e:
                    print(f"Error reading stdout: {e}")
                    pass
            if process.stderr:
                try:
                    data = await process.stderr.read()
                    if data:
                        print(f"Remaining stderr: {data.decode().strip()}")
                except (ValueError, IOError) as e:
                    print(f"Error reading stderr: {e}")
                    pass
                    
        except Exception as e:
            print(f"Exception during cleanup: {e}")
            pass
            
    async def start_bot(self, config_path: Path) -> bool:
        """Start a new bot instance with the given config file.
        
        Args:
            config_path: Path to the config file
            
        Returns:
            bool: True if bot was started successfully
        """
        try:
            # Generate a unique ID for this bot based on the config file name
            bot_id = config_path.stem
            
            # Check if a bot with this ID is already running
            state = self._load_state()
            if bot_id in state:
                try:
                    pid = state[bot_id]['pid']
                    os.kill(pid, 0)  # Check if process exists
                    return False
                except ProcessLookupError:
                    # Process doesn't exist, remove from state
                    pass
            
            # Start the bot process using venv Python
            process = await asyncio.create_subprocess_exec(
                self.venv_python, 'main.py',
                '--config', str(config_path),
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            # Verify process started successfully
            try:
                os.kill(process.pid, 0)
            except ProcessLookupError:
                await self._cleanup_process(process)
                return False
            
            # Start the monitor task if not already running
            if self.monitor_task is None:
                self.monitor_task = asyncio.create_task(self.monitor_processes())
            
            # Monitor the process output for a short time to ensure it starts properly
            try:
                startup_timeout = 5.0  # Give it 5 seconds to start
                start_time = time.time()
                success = False
                
                while time.time() - start_time < startup_timeout:
                    # Check if process has exited
                    if process.returncode is not None:
                        await self._cleanup_process(process)
                        return False
                    
                    # Read any available output during startup
                    try:
                        line = await asyncio.wait_for(process.stderr.readline(), timeout=0.1)
                        if line:
                            line = line.decode().strip()
                            print(f"Bot output: {line}")
                            # Consider the bot started successfully after a short delay
                            # instead of looking for specific output
                            success = True
                    except asyncio.TimeoutError:
                        # No output available, continue monitoring
                        continue
                    except Exception as e:
                        print(f"Error reading bot output: {e}")
                        pass
                
                # Consider the bot started successfully if it's still running after the timeout
                if not success and process.returncode is None:
                    success = True
                
                # Start background task to monitor process output
                async def monitor_output(process, bot_id):
                    try:
                        while True:
                            if process.returncode is not None:
                                break
                            try:
                                line = await process.stderr.readline()
                                if not line:
                                    break
                                # Print the bot's output
                                print(f"Bot {bot_id} output: {line.decode().strip()}")
                            except Exception as e:
                                print(f"Error monitoring bot {bot_id}: {e}")
                                break
                    except Exception as e:
                        print(f"Monitor task exception for bot {bot_id}: {e}")
                        pass
                
                asyncio.create_task(monitor_output(process, bot_id))
                
                # Store the process
                self.bots[bot_id] = process
                self._save_state()
                return True
                
            except Exception:
                await self._cleanup_process(process)
                return False
            
        except Exception:
            return False
    
    async def stop_bot(self, bot_id: str) -> bool:
        """Stop a running bot instance.
        
        Args:
            bot_id: ID of the bot to stop, or "all" to stop all bots
            
        Returns:
            bool: True if bot(s) were stopped successfully
        """
        if bot_id == "all":
            print("Stopping all bots...")
            success = True
            state = self._load_state()
            for bid in list(state.keys()):
                print(f"Stopping bot {bid}...")
                if not await self._stop_single_bot(bid):
                    success = False
            return success
            
        # Stop a single bot
        return await self._stop_single_bot(bot_id)
    
    async def _stop_single_bot(self, bot_id: str) -> bool:
        """Stop a single bot instance.
        
        Args:
            bot_id: ID of the bot to stop
            
        Returns:
            bool: True if bot was stopped successfully
        """
        # Check both local bots and state file
        state = self._load_state()
        process = None
        
        if bot_id in self.bots:
            process = self.bots[bot_id]
        elif bot_id in state:
            try:
                pid = state[bot_id]['pid']
                print(f"Stopping bot {bot_id} with PID {pid}")
                # Try to kill the process
                try:
                    # First try a gentle termination
                    os.kill(pid, signal.SIGTERM)
                    print(f"Sent SIGTERM to process {pid}")
                    
                    # Wait a bit for the process to terminate
                    for i in range(50):  # 5 seconds
                        await asyncio.sleep(0.1)
                        try:
                            os.kill(pid, 0)  # Check if process exists
                        except ProcessLookupError:
                            print(f"Process {pid} terminated successfully")
                            break
                    else:
                        # Process didn't terminate, force kill
                        print(f"Process {pid} didn't terminate, sending SIGKILL")
                        try:
                            os.kill(pid, signal.SIGKILL)
                            print(f"Sent SIGKILL to process {pid}")
                        except ProcessLookupError:
                            print(f"Process {pid} already terminated")
                            pass
                except ProcessLookupError:
                    print(f"Process {pid} not found")
                    pass
                
                # Remove only this bot from state
                if bot_id in state:
                    print(f"Removing {bot_id} from state file")
                    del state[bot_id]
                    with open(self.state_file, 'w') as f:
                        json.dump(state, f)
                    print(f"State file updated, remaining bots: {list(state.keys())}")
                return True
            except Exception as e:
                print(f"Error stopping bot {bot_id}: {e}")
                return False
        else:
            print(f"Bot {bot_id} not found")
            return False
            
        if process:
            try:
                print(f"Cleaning up process for bot {bot_id}")
                await self._cleanup_process(process)
                del self.bots[bot_id]
                
                # Update state file - only remove this bot
                state = self._load_state()
                if bot_id in state:
                    print(f"Removing {bot_id} from state file")
                    del state[bot_id]
                    with open(self.state_file, 'w') as f:
                        json.dump(state, f)
                    print(f"State file updated, remaining bots: {list(state.keys())}")
                return True
            except Exception as e:
                print(f"Error cleaning up process for bot {bot_id}: {e}")
                return False
    
    async def restart_bot(self, bot_id: str) -> bool:
        """Restart a running bot instance.
        
        Args:
            bot_id: ID of the bot to restart
            
        Returns:
            bool: True if bot was restarted successfully
        """
        # Find the config file for this bot
        config_path = self.config_dir / f"{bot_id}.yaml"
        if not config_path.exists():
            return False
        
        # Stop the bot if it's running
        if bot_id in self.bots:
            if not await self.stop_bot(bot_id):
                return False
        
        # Start the bot with the same config
        return await self.start_bot(config_path)
    
    async def quiet_bot(self, bot_id: str) -> bool:
        """Disable song announcements for a running bot.
        
        Args:
            bot_id: ID of the bot to quiet
            
        Returns:
            bool: True if command was sent successfully
        """
        # Check if bot exists in state
        state = self._load_state()
        if bot_id not in state:
            print(f"Bot {bot_id} not found")
            return False
            
        try:
            # Get the socket path for this bot
            socket_path = Path(tempfile.gettempdir()) / f"icecast_bot_{bot_id}.sock"
            if not socket_path.exists():
                print(f"Socket for bot {bot_id} not found at {socket_path}")
                return False
                
            # Send quiet command
            print(f"Sending quiet command to bot {bot_id}")
            reader, writer = await asyncio.open_unix_connection(str(socket_path))
            writer.write(b'quiet')
            await writer.drain()
            writer.close()
            await writer.wait_closed()
            print(f"Quiet command sent to bot {bot_id}")
            return True
        except Exception as e:
            print(f"Error sending quiet command to bot {bot_id}: {e}")
            return False
    
    async def unquiet_bot(self, bot_id: str) -> bool:
        """Enable song announcements for a running bot.
        
        Args:
            bot_id: ID of the bot to unquiet
            
        Returns:
            bool: True if command was sent successfully
        """
        # Check if bot exists in state
        state = self._load_state()
        if bot_id not in state:
            print(f"Bot {bot_id} not found")
            return False
            
        try:
            # Get the socket path for this bot
            socket_path = Path(tempfile.gettempdir()) / f"icecast_bot_{bot_id}.sock"
            if not socket_path.exists():
                print(f"Socket for bot {bot_id} not found at {socket_path}")
                return False
                
            # Send unquiet command
            print(f"Sending unquiet command to bot {bot_id}")
            reader, writer = await asyncio.open_unix_connection(str(socket_path))
            writer.write(b'unquiet')
            await writer.drain()
            writer.close()
            await writer.wait_closed()
            print(f"Unquiet command sent to bot {bot_id}")
            return True
        except Exception as e:
            print(f"Error sending unquiet command to bot {bot_id}: {e}")
            return False
    
    async def list_bots(self) -> bool:
        """List all running bots.
        
        Returns:
            bool: True if any bots are running
        """
        state = self._load_state()
        
        # Check if any bots are running
        if not state:
            print("No bots running")
            return False
            
        # Track unique PIDs to avoid duplicates
        seen_pids = set()
        unique_bots = {}
        
        # Filter out duplicates based on PID
        for bot_id, info in state.items():
            pid = info.get('pid')
            if pid and pid not in seen_pids:
                seen_pids.add(pid)
                unique_bots[bot_id] = info
        
        if not unique_bots:
            print("No bots running")
            return False
            
        # Print header
        print("\nRunning Bots:")
        print("-" * 80)
        print(f"{'ID':<20} {'PID':<8} {'Status':<10} {'Command':<40}")
        print("-" * 80)
        
        # Print each bot's status
        for bot_id, info in unique_bots.items():
            pid = info.get('pid')
            config = info.get('config', '')
            
            # Check if process is still running
            try:
                os.kill(pid, 0)
                status = "running"
            except ProcessLookupError:
                status = "stopped"
                continue  # Skip stopped processes
            except Exception:
                status = "unknown"
                
            # Get command line
            try:
                with open(f"/proc/{pid}/cmdline", 'rb') as f:
                    cmdline = f.read().replace(b'\0', b' ').decode()
            except Exception:
                cmdline = f"Unknown (PID: {pid})"
                
            print(f"{bot_id:<20} {pid:<8} {status:<10} {cmdline:<40}")
            
        print("-" * 80)
        print()
        
        return True
    
    async def cleanup(self):
        """Clean up all running bots."""
        if self.monitor_task:
            self.monitor_task.cancel()
            try:
                await self.monitor_task
            except asyncio.CancelledError:
                pass
        
        await self.stop_bot("all")

async def main():
    parser = argparse.ArgumentParser(description='Icecast IRC Bot Manager')
    parser.add_argument('--config', type=str, help='Path to config file or directory')
    parser.add_argument('command', choices=['start', 'stop', 'restart', 'list', 'quiet', 'unquiet'], help='Command to execute')
    parser.add_argument('bot_id', nargs='?', help='Bot ID for start/stop/restart/quiet/unquiet commands, or "all" to stop all bots')

    args = parser.parse_args()
    
    manager = BotManager()
    should_cleanup = False  # Only cleanup for certain commands
    
    try:
        if args.command == 'list':
            if await manager.list_bots():
                # If list_bots returns True, we've printed the list
                pass
        elif args.command == 'start':
            should_cleanup = True  # Need cleanup for start command
            if not args.config:
                print("Error: --config required for start command")
                sys.exit(1)
            config_path = Path(args.config)
            if config_path.is_dir():
                # Start all bots in directory
                success = True
                for config_file in config_path.glob('*.yaml'):
                    if not await manager.start_bot(config_file):
                        success = False
                if not success:
                    sys.exit(1)
            else:
                # Start single bot
                if not await manager.start_bot(config_path):
                    sys.exit(1)
                
            # If we started any bots successfully, keep running until interrupted
            if manager.bots:
                try:
                    # Keep the manager running
                    while True:
                        await asyncio.sleep(1)
                        if not manager.bots:
                            break
                except asyncio.CancelledError:
                    pass
                
        elif args.command == 'stop':
            # Don't need cleanup for stop command as it already cleans up
            should_cleanup = False
            if not args.bot_id:
                print("Error: bot_id required for stop command (use 'all' to stop all bots)")
                sys.exit(1)
            if not await manager.stop_bot(args.bot_id):
                sys.exit(1)
            
        elif args.command == 'restart':
            should_cleanup = True  # Need cleanup for restart command
            if not args.bot_id:
                print("Error: bot_id required for restart command")
                sys.exit(1)
            if args.bot_id == "all":
                print("Error: restart all is not supported")
                sys.exit(1)
            if not await manager.restart_bot(args.bot_id):
                sys.exit(1)
            
            # If we restarted successfully, keep running until interrupted
            if manager.bots:
                try:
                    # Keep the manager running
                    while True:
                        await asyncio.sleep(1)
                        if not manager.bots:
                            break
                except asyncio.CancelledError:
                    pass
        
        elif args.command == 'quiet':
            should_cleanup = False  # Don't need cleanup for quiet command
            if not args.bot_id:
                print("Error: bot_id required for quiet command")
                sys.exit(1)
            if args.bot_id == "all":
                print("Error: quiet all is not supported")
                sys.exit(1)
            if not await manager.quiet_bot(args.bot_id):
                sys.exit(1)
                
        elif args.command == 'unquiet':
            should_cleanup = False  # Don't need cleanup for unquiet command
            if not args.bot_id:
                print("Error: bot_id required for unquiet command")
                sys.exit(1)
            if args.bot_id == "all":
                print("Error: unquiet all is not supported")
                sys.exit(1)
            if not await manager.unquiet_bot(args.bot_id):
                sys.exit(1)
            
    except KeyboardInterrupt:
        should_cleanup = True  # Need cleanup for keyboard interrupt
    except Exception:
        should_cleanup = True  # Need cleanup for errors
        sys.exit(1)
    finally:
        # Only clean up if we need to
        if should_cleanup:
            await manager.cleanup()

if __name__ == '__main__':
    # Set up signal handlers for graceful shutdown
    def signal_handler(sig, frame):
        # Don't exit immediately, let the cleanup happen
        asyncio.get_event_loop().stop()
    
    for sig in (signal.SIGTERM, signal.SIGINT):
        signal.signal(sig, signal_handler)
    
    # Run the manager
    asyncio.run(main())