#!/usr/bin/env python3
"""
Enable multiple displays on a SPICE VM by sending monitor configuration.
Uses SpiceMainChannel.set_display() to configure displays directly.
"""

import gi
import sys

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):
        self.uri = uri
        self.num_displays = num_displays
        self.width = width
        self.height = height
        self.session = None
        self.main_channel = None
        self.display_channels = []
        self.loop = GLib.MainLoop()
        self.configured = 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')
        print(f"New channel: type={channel_type}, id={channel_id}")

        if channel_type == CHANNEL_MAIN:
            self.main_channel = channel
            channel.connect_after('channel-event', self.on_channel_event)
        elif channel_type == CHANNEL_DISPLAY:
            self.display_channels.append((channel_id, channel))
            print(f"  Display channel {channel_id} added")

    def on_channel_event(self, channel, event):
        """Handle channel events"""
        print(f"Channel event: {event}")
        if event == SpiceClientGLib.ChannelEvent.OPENED:
            print("Main channel opened, configuring displays...")
            GLib.timeout_add(2000, self.configure_monitors)

    def configure_monitors(self):
        """Configure multiple monitors via SPICE protocol"""
        print(f"\n=== Configuring {self.num_displays} displays ===")
        print(f"Display channels available: {len(self.display_channels)}")

        if not self.main_channel:
            print("No main channel!")
            GLib.timeout_add(1000, self.quit)
            return False

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

            print(f"Setting display {i}: {self.width}x{self.height} at ({x}, {y})")

            try:
                # Enable the display
                self.main_channel.set_display_enabled(i, True)

                # Set display geometry (id, x, y, width, height)
                self.main_channel.set_display(i, x, y, self.width, self.height)
            except Exception as e:
                print(f"  Error setting display {i}: {e}")

        # Send the configuration immediately
        print("\nSending monitor config to guest...")
        try:
            self.main_channel.send_monitor_config()
            self.configured = True
            print("Monitor config sent!")
        except Exception as e:
            print(f"Error sending config: {e}")

        # Wait a bit then check agent status and quit
        GLib.timeout_add(3000, self.check_and_quit)
        return False

    def check_and_quit(self):
        """Check final status and quit"""
        if self.main_channel:
            agent_connected = self.main_channel.get_property('agent-connected')
            print(f"\nAgent connected: {agent_connected}")
        self.quit()
        return False

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

    def run(self):
        print(f"Connecting to {self.uri}...")
        print(f"Target: {self.num_displays} displays at {self.width}x{self.height}")

        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")
            return False

        # Fallback timeout to configure monitors
        GLib.timeout_add(5000, self.configure_monitors)
        GLib.timeout_add(15000, self.quit)
        self.loop.run()

        print(f"\n=== Result ===")
        print(f"Configured: {self.configured}")
        print(f"Display channels: {len(self.display_channels)}")
        return self.configured

if __name__ == '__main__':
    uri = sys.argv[1] if len(sys.argv) > 1 else 'spice://localhost:5930'
    num_displays = int(sys.argv[2]) if len(sys.argv) > 2 else 3

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