#!/usr/bin/env python3
"""
Enable multiple displays on a SPICE VM by sending monitor configuration.
Retries until the SPICE agent in the guest is connected.
"""

import gi
import sys
import time

gi.require_version('SpiceClientGLib', '2.0')
from gi.repository import SpiceClientGLib, GLib

# Channel types (from spice-protocol)
CHANNEL_MAIN = 1
CHANNEL_DISPLAY = 2

class SpiceDisplayEnabler:
    def __init__(self, uri, num_displays=3, width=1920, height=1080, timeout=60):
        self.uri = uri
        self.num_displays = num_displays
        self.width = width
        self.height = height
        self.timeout = timeout
        self.session = None
        self.main_channel = None
        self.display_channels = []
        self.loop = GLib.MainLoop()
        self.configured = False
        self.agent_connected = False
        self.config_sent = False

    def on_channel_new(self, session, channel):
        """Handle new channel creation"""
        channel_type = channel.get_property('channel-type')
        channel_id = channel.get_property('channel-id')

        if channel_type == CHANNEL_MAIN:
            self.main_channel = channel
            channel.connect_after('channel-event', self.on_channel_event)
            # Check agent status periodically
            GLib.timeout_add(500, self.check_agent_and_configure)
        elif channel_type == CHANNEL_DISPLAY:
            self.display_channels.append((channel_id, channel))

    def on_channel_event(self, channel, event):
        """Handle channel events"""
        if event == SpiceClientGLib.ChannelEvent.OPENED:
            # Start checking for agent
            GLib.timeout_add(100, self.check_agent_and_configure)

    def check_agent_and_configure(self):
        """Check if agent is connected and configure if ready"""
        if self.config_sent:
            return False  # Stop checking

        if not self.main_channel:
            return True  # Keep checking

        self.agent_connected = self.main_channel.get_property('agent-connected')

        if self.agent_connected and not self.config_sent:
            print(f"Agent connected! Configuring {self.num_displays} displays...")
            self.configure_monitors()
            return False  # Stop checking

        return True  # Keep checking

    def configure_monitors(self):
        """Configure multiple monitors via SPICE protocol"""
        if self.config_sent:
            return

        if not self.main_channel:
            print("No main channel!")
            return

        # Enable and configure each display
        for i in range(self.num_displays):
            x = i * self.width  # Position displays side by side
            y = 0

            try:
                self.main_channel.update_display_enabled(i, True, False)
                self.main_channel.update_display(i, x, y, self.width, self.height, False)
            except Exception as e:
                print(f"  Error setting display {i}: {e}")

        # Send the configuration
        try:
            self.main_channel.send_monitor_config()
            self.config_sent = True
            self.configured = True
            print(f"Configured {self.num_displays} displays at {self.width}x{self.height}")
        except Exception as e:
            print(f"Error sending config: {e}")

        # Quit after a short delay
        GLib.timeout_add(1000, self.quit)

    def quit(self):
        self.loop.quit()
        return False

    def on_timeout(self):
        """Handle overall timeout"""
        if not self.configured:
            print(f"Timeout after {self.timeout}s - agent not connected")
        self.quit()
        return False

    def run(self):
        print(f"Connecting to {self.uri}...")
        print(f"Waiting up to {self.timeout}s for agent...")

        self.session = SpiceClientGLib.Session()
        self.session.set_property('uri', self.uri)
        self.session.connect_after('channel-new', self.on_channel_new)

        if not self.session.connect():
            print("Failed to connect to SPICE server")
            return False

        # Set overall timeout
        GLib.timeout_add(self.timeout * 1000, self.on_timeout)

        self.loop.run()

        if self.configured:
            print(f"Success: {self.num_displays} displays enabled")
        else:
            print("Failed: Could not enable displays")

        return self.configured


def main():
    import argparse
    parser = argparse.ArgumentParser(description='Enable SPICE VM displays')
    parser.add_argument('uri', nargs='?', default='spice://localhost:5930',
                        help='SPICE URI (default: spice://localhost:5930)')
    parser.add_argument('num_displays', nargs='?', type=int, default=3,
                        help='Number of displays to enable (default: 3)')
    parser.add_argument('--timeout', '-t', type=int, default=60,
                        help='Timeout in seconds (default: 60)')
    parser.add_argument('--width', '-W', type=int, default=1920,
                        help='Display width (default: 1920)')
    parser.add_argument('--height', '-H', type=int, default=1080,
                        help='Display height (default: 1080)')

    args = parser.parse_args()

    enabler = SpiceDisplayEnabler(
        args.uri,
        args.num_displays,
        args.width,
        args.height,
        args.timeout
    )
    success = enabler.run()
    sys.exit(0 if success else 1)


if __name__ == '__main__':
    main()