import network
import rp2
import time
import sys
import machine
from machine import Pin
import select
import socket

pico_led = Pin("LED", Pin.OUT)

SCAN_TIMEOUT = 5000

class WiFi2USB:
    def __init__(self):
        self.usb_net = network.USB_NET()
        self.usb_net.active(True)

        self.wlan = network.WLAN(network.STA_IF)
        self.wlan.active(True)

        self.networks = []
        self.current_ssid = ""
        self.current_password = ""
        self.connected = False

    def scan_networks(self):
        """Scan for available WiFi networks"""
        pico_led.on()
        print("Scanning for networks...")
        self.networks = self.wlan.scan()
        pico_led.off()
        return self.networks

    def connect_wifi(self, ssid, password):
        """Connect to a WiFi network with given SSID and password"""
        pico_led.off()
        print(f"Connecting to {ssid}...")

        self.wlan.connect(ssid, password)

        start_time = time.ticks_ms()
        while not self.wlan.isconnected():
            # Check for cancel with BOOTSEL button
            if rp2.bootsel_button() == 1:
                print("\nConnection attempt cancelled")
                return False

            if time.ticks_diff(time.ticks_ms(), start_time) > SCAN_TIMEOUT:
                print("\nConnection timeout - please try again")
                return False

            pico_led.on()
            time.sleep(0.1)
            pico_led.off()
            time.sleep(0.1)

        ip = self.wlan.ifconfig()[0]
        print(f'Connected on {ip}')
        pico_led.on()
        self.connected = True
        self.current_ssid = ssid
        self.current_password = password
        return True

    def disconnect_wifi(self):
        """Disconnect from WiFi network"""
        if self.connected:
            self.wlan.disconnect()
            self.connected = False
            pico_led.off()
            print("Disconnected from WiFi")

    def get_network_info(self):
        """Return current network configuration"""
        if self.connected:
            return {
                'ssid': self.current_ssid,
                'ip': self.wlan.ifconfig()[0],
                'subnet': self.wlan.ifconfig()[1],
                'gateway': self.wlan.ifconfig()[2],
                'dns': self.wlan.ifconfig()[3],
                'status': 'Connected'
            }
        else:
            return {'status': 'Not connected'}

    def test_connectivity(self):
        """Test network connectivity"""
        if not self.connected:
            print("Not connected to any network")
            return False

        try:
            ip = self.wlan.ifconfig()[0]
            print(f"Connected with IP: {ip}")
            print("Gateway: " + self.wlan.ifconfig()[2])

            # TODO: ping a reliable site here
            print("Connectivity test passed")
            return True
        except:
            print("Connectivity test failed")
            return False

    def get_usb_status(self):
        """Return USB RNDIS status"""
        return {
            'active': self.usb_net.active(),
            'ifconfig': self.usb_net.ifconfig()
        }

    def forward_traffic(self):
        """Forward traffic between WiFi and USB interfaces"""
        try:
            wlan_sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_IP)
            usb_sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_IP)

            wlan_sock.bind((self.wlan.ifconfig()[0], 0))
            usb_sock.bind((self.usb_net.ifconfig()[0], 0))

            poller = select.poll()
            poller.register(wlan_sock, select.POLLIN)
            poller.register(usb_sock, select.POLLIN)

            print("Forwarding traffic between WiFi and USB interfaces...")

            while True:
                events = poller.poll(1000)
                for fd, event in events:
                    if fd == wlan_sock.fileno():
                        data, addr = wlan_sock.recvfrom(1500)
                        usb_sock.sendto(data, (self.usb_net.ifconfig()[0], 0))
                    elif fd == usb_sock.fileno():
                        data, addr = usb_sock.recvfrom(1500)
                        wlan_sock.sendto(data, (self.wlan.ifconfig()[0], 0))
        except OSError as e:
            print(f"Socket error: {e}")
        except Exception as e:
            print(f"Unexpected error: {e}")

class CLI:
    def __init__(self, wifi2usb):
        self.wifi2usb = wifi2usb
        self.current_menu = self.main_menu

    def start(self):
        """Start the CLI interface"""
        print("\n" + "=" * 50)
        print("WiFi2USB Dongle - CLI Interface")
        print("=" * 50)

        while True:
            try:
                self.current_menu()
            except KeyboardInterrupt:
                print("\nExiting CLI Interface")
                break

    def main_menu(self):
        """Display the main menu"""
        print("\nMAIN MENU:")
        print("1. Reconfigure network")
        print("2. See network config")
        print("3. Test network connectivity")
        print("4. Enter microPython REPL")
        print("0. Exit")

        choice = input("\nEnter your choice: ")

        if choice == "1":
            self.current_menu = self.network_selection_menu
        elif choice == "2":
            self.network_config_menu()
        elif choice == "3":
            self.connectivity_test_menu()
        elif choice == "4":
            self.enter_repl()
        elif choice == "0":
            print("Exiting...")
            sys.exit()
        else:
            print("Invalid choice, please try again")

    def network_selection_menu(self):
        """Display network selection menu"""
        print("\nScanning for networks...")
        networks = self.wifi2usb.scan_networks()

        print("\nAvailable Networks:")
        print("   Network Name                [#Channel] |Other info|")
        print("-" * 60)

        for i, network in enumerate(networks, 1):
            try:
                ssid = network[0].decode('utf-8')
                bssid = network[1]  # MAC address in bytes
                channel = network[2]
                rssi = network[3]
                authmode = network[4]

                # Format MAC address
                mac = ':'.join('{:02x}'.format(b) for b in bssid)

                # Safe access to security types with fallback
                security_types = ["Open", "WEP", "WPA-PSK", "WPA2-PSK", "WPA/WPA2-PSK", "WPA3"]
                if 0 <= authmode < len(security_types):
                    security = security_types[authmode]
                else:
                    security = f"Unknown({authmode})"

                print(f"{i}. {ssid:<25} [#{channel}]     |mac:{mac}, db:{rssi}, {security}|")
            except Exception as e:
                print(f"{i}. [Error reading network info: {e}]")

        print("-" * 60)
        print("0: go back to last menu")

        try:
            choice = input("\nSelect network: ")

            if choice == "0":
                self.current_menu = self.main_menu
                return

            idx = int(choice) - 1
            if 0 <= idx < len(networks):
                selected_ssid = networks[idx][0].decode('utf-8')
                password = input(f"Enter password for {selected_ssid}: ")

                if self.wifi2usb.connect_wifi(selected_ssid, password):
                    print(f"Successfully connected to {selected_ssid}")
                    self.wifi2usb.forward_traffic()
                else:
                    print(f"Failed to connect to {selected_ssid}")

                # Return to main menu after connection attempt
                self.current_menu = self.main_menu
            else:
                print("Invalid selection")
        except ValueError:
            print("Please enter a valid number")

    def network_config_menu(self):
        """Display current network configuration"""
        print("\nCurrent Network Configuration:")
        print("-" * 50)

        # WiFi status
        wifi_info = self.wifi2usb.get_network_info()
        print("WiFi Status:", wifi_info['status'])

        if wifi_info['status'] == 'Connected':
            print(f"SSID: {wifi_info['ssid']}")
            print(f"IP Address: {wifi_info['ip']}")
            print(f"Subnet Mask: {wifi_info['subnet']}")
            print(f"Gateway: {wifi_info['gateway']}")
            print(f"DNS Server: {wifi_info['dns']}")

        # USB RNDIS status
        usb_status = self.wifi2usb.get_usb_status()
        print("\nUSB Network Interface:")
        print(f"Active: {usb_status['active']}")
        print(f"IP Configuration: {usb_status['ifconfig']}")

        print("-" * 50)
        input("Press Enter to return to main menu...")

    def connectivity_test_menu(self):
        """Test network connectivity"""
        print("\nTesting Network Connectivity:")
        print("-" * 50)

        self.wifi2usb.test_connectivity()

        print("-" * 50)
        input("Press Enter to return to main menu...")

    def enter_repl(self):
        """Enter MicroPython REPL"""
        print("\nEntering MicroPython REPL...")
        print('Type "exit()" to return to CLI')
        print("-" * 50)

        # TODO: Replace this, it sucks
        while True:
            try:
                code = input(">>> ")
                if code.strip() == "exit()":
                    break
                try:
                    result = eval(code)
                    if result is not None:
                        print(result)
                except SyntaxError:
                    exec(code)
            except Exception as e:
                print(f"Error: {e}")

        print("-" * 50)
        print("Returned from REPL")

def main():
    wifi2usb = WiFi2USB()
    cli = CLI(wifi2usb)
    cli.start()

if __name__ == "__main__":
    main()