#!/usr/bin/env python3
"""
Map ESP32 USB ports to IP addresses
Creates and manages mapping between /dev/ttyUSB* and assigned IPs
"""

import serial.tools.list_ports
import argparse
import json
import glob
import re
from pathlib import Path

class USBIPMapper:
    def __init__(self, start_ip="192.168.1.51", config_file="usb_ip_map.json"):
        self.start_ip = start_ip
        self.config_file = config_file
        self.mapping = {}
        
    def get_ip_for_index(self, index):
        """Calculate IP address for a given index"""
        ip_parts = self.start_ip.split('.')
        base_ip = int(ip_parts[3])
        ip_parts[3] = str(base_ip + index)
        return '.'.join(ip_parts)
    
    def extract_usb_number(self, port):
        """Extract number from /dev/ttyUSBX"""
        match = re.search(r'ttyUSB(\d+)', port)
        if match:
            return int(match.group(1))
        return None
    
    def detect_devices(self):
        """Detect all ESP32 USB devices and create mapping"""
        devices = sorted(glob.glob('/dev/ttyUSB*'))
        
        print(f"\n{'='*70}")
        print(f"ESP32 USB to IP Address Mapping")
        print(f"{'='*70}")
        print(f"Start IP: {self.start_ip}")
        print(f"Detected {len(devices)} USB device(s)\n")
        
        self.mapping = {}
        
        for idx, port in enumerate(devices):
            usb_num = self.extract_usb_number(port)
            ip = self.get_ip_for_index(idx)
            
            # Get device info
            try:
                ports = serial.tools.list_ports.comports()
                device_info = next((p for p in ports if p.device == port), None)
                if device_info:
                    serial_num = device_info.serial_number or "Unknown"
                    description = device_info.description or "Unknown"
                else:
                    serial_num = "Unknown"
                    description = "Unknown"
            except:
                serial_num = "Unknown"
                description = "Unknown"
            
            self.mapping[port] = {
                'index': idx,
                'usb_number': usb_num,
                'ip': ip,
                'serial': serial_num,
                'description': description
            }
            
            print(f"[{idx:2d}] {port:14s} → {ip:15s} (USB{usb_num}, SN: {serial_num})")
        
        print(f"\n{'='*70}")
        print(f"Total: {len(devices)} devices mapped")
        print(f"IP Range: {self.mapping[devices[0]]['ip']} - {self.mapping[devices[-1]]['ip']}" if devices else "")
        print(f"{'='*70}\n")
        
        return self.mapping
    
    def save_mapping(self):
        """Save mapping to JSON file"""
        with open(self.config_file, 'w') as f:
            json.dump(self.mapping, f, indent=2)
        print(f"✓ Mapping saved to {self.config_file}")
    
    def load_mapping(self):
        """Load mapping from JSON file"""
        try:
            with open(self.config_file, 'r') as f:
                self.mapping = json.load(f)
            print(f"✓ Mapping loaded from {self.config_file}")
            return self.mapping
        except FileNotFoundError:
            print(f"✗ No saved mapping found at {self.config_file}")
            return {}
    
    def get_ip(self, port):
        """Get IP address for a specific USB port"""
        if port in self.mapping:
            return self.mapping[port]['ip']
        return None
    
    def get_port(self, ip):
        """Get USB port for a specific IP address"""
        for port, info in self.mapping.items():
            if info['ip'] == ip:
                return port
        return None
    
    def print_mapping(self):
        """Print current mapping"""
        if not self.mapping:
            print("No mapping loaded. Run with --detect first.")
            return
        
        print(f"\n{'='*70}")
        print(f"Current USB to IP Mapping")
        print(f"{'='*70}")
        for port, info in sorted(self.mapping.items(), key=lambda x: x[1]['index']):
            print(f"[{info['index']:2d}] {port:14s} → {info['ip']:15s} (USB{info['usb_number']})")
        print(f"{'='*70}\n")
    
    def export_bash_script(self, filename="usb_ip_vars.sh"):
        """Export mapping as bash variables"""
        with open(filename, 'w') as f:
            f.write("#!/bin/bash\n")
            f.write("# USB to IP mapping - Auto-generated\n\n")
            
            # Create associative array
            f.write("declare -A USB_TO_IP\n")
            for port, info in self.mapping.items():
                f.write(f"USB_TO_IP[{port}]=\"{info['ip']}\"\n")
            
            f.write("\n# Create reverse mapping\n")
            f.write("declare -A IP_TO_USB\n")
            for port, info in self.mapping.items():
                f.write(f"IP_TO_USB[{info['ip']}]=\"{port}\"\n")
            
            f.write("\n# Helper functions\n")
            f.write("get_ip_for_usb() { echo \"${USB_TO_IP[$1]}\"; }\n")
            f.write("get_usb_for_ip() { echo \"${IP_TO_USB[$1]}\"; }\n")
        
        print(f"✓ Bash script exported to {filename}")
        print(f"  Usage: source {filename} && get_ip_for_usb /dev/ttyUSB0")

def main():
    parser = argparse.ArgumentParser(
        description='Map ESP32 USB ports to IP addresses',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Detect devices and create mapping
  %(prog)s --detect
  
  # Detect and save to file
  %(prog)s --detect --save
  
  # Load saved mapping and display
  %(prog)s --load --print
  
  # Get IP for specific USB port
  %(prog)s --load --port /dev/ttyUSB5
  
  # Get USB port for specific IP
  %(prog)s --load --ip 192.168.1.55
  
  # Export as bash script
  %(prog)s --load --export
  
  # Use custom IP range
  %(prog)s --detect --start-ip 10.0.0.100
        """
    )
    
    parser.add_argument('--detect', action='store_true',
                        help='Detect USB devices and create mapping')
    parser.add_argument('--save', action='store_true',
                        help='Save mapping to file')
    parser.add_argument('--load', action='store_true',
                        help='Load mapping from file')
    parser.add_argument('--print', action='store_true',
                        help='Print current mapping')
    parser.add_argument('--start-ip', default='192.168.1.51',
                        help='Starting IP address (default: 192.168.1.51)')
    parser.add_argument('--config', default='usb_ip_map.json',
                        help='Config file path (default: usb_ip_map.json)')
    parser.add_argument('--port', metavar='PORT',
                        help='Get IP for specific USB port (e.g., /dev/ttyUSB5)')
    parser.add_argument('--ip', metavar='IP',
                        help='Get USB port for specific IP address')
    parser.add_argument('--export', action='store_true',
                        help='Export mapping as bash script')
    
    args = parser.parse_args()
    
    mapper = USBIPMapper(start_ip=args.start_ip, config_file=args.config)
    
    # Detect devices
    if args.detect:
        mapper.detect_devices()
        if args.save:
            mapper.save_mapping()
    
    # Load mapping
    if args.load:
        mapper.load_mapping()
    
    # Print mapping
    if args.print:
        mapper.print_mapping()
    
    # Query specific port
    if args.port:
        if not mapper.mapping:
            mapper.load_mapping()
        ip = mapper.get_ip(args.port)
        if ip:
            print(f"{args.port} → {ip}")
        else:
            print(f"Port {args.port} not found in mapping")
    
    # Query specific IP
    if args.ip:
        if not mapper.mapping:
            mapper.load_mapping()
        port = mapper.get_port(args.ip)
        if port:
            print(f"{args.ip} → {port}")
        else:
            print(f"IP {args.ip} not found in mapping")
    
    # Export bash script
    if args.export:
        if not mapper.mapping:
            mapper.load_mapping()
        mapper.export_bash_script()
    
    # Default: detect and print
    if not any([args.detect, args.load, args.print, args.port, args.ip, args.export]):
        mapper.detect_devices()

if __name__ == '__main__':
    main()