import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, GLib, Gdk
import threading
import subprocess
import os
import json

class RemotePairingDialog(Gtk.Dialog):
    def __init__(self, parent):
        super().__init__(title="Axiom Zero — Remote Server Pairing", transient_for=parent, flags=0)
        self.set_default_size(450, 250)
        self.set_border_width(20)
        self.set_modal(True)
        
        provider = Gtk.CssProvider()
        provider.load_from_data(b"""
            label.title { font-size: 16px; font-weight: bold; color: #00ffcc; margin-bottom: 10px; }
            entry { background: #1a1a1a; color: #fff; padding: 8px; border: 1px solid #333; }
            button { background: #004d40; color: #fff; font-weight: bold; padding: 10px; }
            button:hover { background: #00796b; }
        """)
        Gtk.StyleContext.add_provider_for_screen(
            Gdk.Screen.get_default(), provider, Gtk.STYLE_PROVIDER_PRIORITY_APPLICATION
        )
        
        box = self.get_content_area()
        box.set_spacing(10)
        
        lbl_title = Gtk.Label(label="Connect to Remote Daemon", xalign=0)
        lbl_title.get_style_context().add_class("title")
        box.pack_start(lbl_title, False, False, 0)
        
        box.pack_start(Gtk.Label(label="Server IP or Hostname:", xalign=0), False, False, 0)
        self.entry_ip = Gtk.Entry()
        self.entry_ip.set_placeholder_text("e.g. 192.168.1.50")
        box.pack_start(self.entry_ip, False, False, 0)
        
        box.pack_start(Gtk.Label(label="SSH Username:", xalign=0), False, False, 0)
        self.entry_user = Gtk.Entry()
        self.entry_user.set_text("root")
        box.pack_start(self.entry_user, False, False, 0)
        
        self.lbl_status = Gtk.Label(label="", xalign=0)
        self.lbl_status.override_color(Gtk.StateFlags.NORMAL, Gdk.RGBA(1.0, 0.4, 0.4, 1.0))
        box.pack_start(self.lbl_status, False, False, 0)
        
        self.btn_connect = Gtk.Button(label="Pair & Connect")
        self.btn_connect.connect("clicked", self.on_connect_clicked)
        box.pack_start(self.btn_connect, False, False, 10)
        
        self.show_all()
        
    def on_connect_clicked(self, widget):
        ip = self.entry_ip.get_text().strip()
        user = self.entry_user.get_text().strip()
        if not ip or not user:
            self.lbl_status.set_text("Please enter Server IP and Username.")
            return
            
        self.btn_connect.set_sensitive(False)
        self.lbl_status.override_color(Gtk.StateFlags.NORMAL, Gdk.RGBA(0.6, 0.8, 1.0, 1.0))
        self.lbl_status.set_text("Connecting via SSH to retrieve TLS certificates...")
        
        threading.Thread(target=self.do_pairing, args=(ip, user), daemon=True).start()
        
    def do_pairing(self, ip, user):
        axiom_dir = os.path.expanduser("~/.axiom-zero")
        server_dir = os.path.join(axiom_dir, ip)
        os.makedirs(server_dir, exist_ok=True, mode=0o700)
        
        certs = [("client.pem", "client.pem"), ("client.key", "client.key"), ("ca.pem", "ca.pem")]
        success = True
        
        for remote_file, local_file in certs:
            # We assume key-based SSH authentication is set up.
            cmd = ["scp", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5",
                   f"{user}@{ip}:/var/lib/axiom-zero/{remote_file}", f"{server_dir}/{local_file}"]
            try:
                res = subprocess.run(cmd, capture_output=True, text=True)
                if res.returncode != 0:
                    success = False
                    error_msg = res.stderr.strip() or "SSH connection failed"
                    break
                os.chmod(f"{server_dir}/{local_file}", 0o600)
            except Exception as e:
                success = False
                error_msg = str(e)
                break
                
        if success:
            cfg = {
                "server_host": ip,
                "server_port": 7443,
                "tls_verify": True,
                "ca_cert": f"{server_dir}/ca.pem",
                "client_cert": f"{server_dir}/client.pem",
                "client_key": f"{server_dir}/client.key"
            }
            conn_file = f"{axiom_dir}/connection.json"
            configs = []
            if os.path.exists(conn_file):
                try:
                    with open(conn_file, "r") as f:
                        data = json.load(f)
                        configs = data if isinstance(data, list) else [data]
                except Exception:
                    pass
            
            # Remove existing config for this IP if present, then append new
            configs = [c for c in configs if c.get("server_host") != ip]
            configs.append(cfg)
            
            with open(conn_file, "w") as f:
                json.dump(configs, f, indent=2)
            GLib.idle_add(self.finish_pairing, True, "", cfg)
        else:
            GLib.idle_add(self.finish_pairing, False, error_msg, None)
            
    def finish_pairing(self, success, error_msg, new_cfg):
        if success:
            # We can attach the new config to the dialog object so the caller can read it
            self.new_cfg = new_cfg
            self.response(Gtk.ResponseType.OK)
        else:
            self.btn_connect.set_sensitive(True)
            self.lbl_status.override_color(Gtk.StateFlags.NORMAL, Gdk.RGBA(1.0, 0.4, 0.4, 1.0))
            self.lbl_status.set_text(f"Pairing Failed: {error_msg}")