# ssh_client.py - Netmiko SSH connections to FS switches
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from netmiko import ConnectHandler, NetmikoTimeoutException, NetmikoAuthenticationException
from config import SSH_USERNAME, SSH_PASSWORD, SSH_PORT, SSH_TIMEOUT, DEVICE_TYPE
from parser import parse_lldp_neighbors, parse_mgmt_ip_from_interfaces

logger = logging.getLogger(__name__)


def connect_and_query(ip):
    """
    SSH into a switch, run LLDP + interface commands.
    Returns dict with switch info and neighbors, or error info.
    """
    device = {
        "device_type": DEVICE_TYPE,
        "host": ip,
        "username": SSH_USERNAME,
        "password": SSH_PASSWORD,
        "port": SSH_PORT,
        "timeout": SSH_TIMEOUT,
        "global_delay_factor": 2,
        "fast_cli": False,
        "conn_timeout": SSH_TIMEOUT,
#        "ssh_config_file": None,
#        "disabled_algorithms": {},
#        "transport": "paramiko",
    }

    # FS switches use legacy ssh-rsa keys — must be explicitly allowed
    import paramiko
    _orig_connect = ConnectHandler.__init__

    try:
        logger.info(f"Connecting to {ip}...")

        # Patch paramiko transport to allow legacy key algorithms
        import paramiko.transport as _pt
        _orig_preferred_keys = _pt.Transport._preferred_keys
        _pt.Transport._preferred_keys = (
            "ssh-rsa", "rsa-sha2-256", "rsa-sha2-512",
            "ecdsa-sha2-nistp256", "ecdsa-sha2-nistp384",
            "ecdsa-sha2-nistp521", "ssh-ed25519",
        )

        with ConnectHandler(**device) as conn:
            # Restore after connect
            _pt.Transport._preferred_keys = _orig_preferred_keys

            # Get hostname from prompt
            hostname = conn.find_prompt().replace('#', '').replace('>', '').strip()

            # Get LLDP neighbors
            lldp_output = conn.send_command("show lldp neighbors", read_timeout=30)

            # Get management IP from interface brief
            intf_output = conn.send_command("show ip interface brief", read_timeout=30)
            mgmt_ip = parse_mgmt_ip_from_interfaces(intf_output) or ip

            # Get chassis ID from LLDP local info
            local_info_output = conn.send_command("show lldp local-information", read_timeout=30)
            chassis_id = _extract_local_chassis(local_info_output) or ip

            # Parse neighbors
            neighbors = parse_lldp_neighbors(lldp_output, chassis_id, hostname, mgmt_ip)

            # Get system description
            sys_desc = _extract_system_desc(local_info_output)

            logger.info(f"  {hostname} ({ip}): {len(neighbors)} neighbors found")

            return {
                "success": True,
                "ip": ip,
                "hostname": hostname,
                "chassis_id": chassis_id,
                "mgmt_ip": mgmt_ip,
                "description": sys_desc,
                "neighbors": neighbors,
            }

    except NetmikoAuthenticationException:
        logger.error(f"Auth failed for {ip}")
        return {"success": False, "ip": ip, "error": "Authentication failed"}
    except NetmikoTimeoutException:
        logger.error(f"Timeout connecting to {ip}")
        return {"success": False, "ip": ip, "error": "Connection timed out"}
    except Exception as e:
        logger.error(f"Error on {ip}: {e}")
        return {"success": False, "ip": ip, "error": str(e)}


def _extract_local_chassis(output):
    import re
    m = re.search(r'Chassis ID\s*:\s*([0-9A-Fa-f\-:]+)', output)
    return m.group(1).strip() if m else None


def _extract_system_desc(output):
    import re
    m = re.search(r'System Description\s*:\s*(.+)', output)
    return m.group(1).strip() if m else ''


def scan_all_switches(ip_list, progress_callback=None, max_workers=10):
    """
    Scan all switches in parallel.
    progress_callback(done, total, current_ip, result) called after each switch.
    Returns list of results.
    """
    results = []
    total = len(ip_list)
    done = 0

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_ip = {executor.submit(connect_and_query, ip): ip for ip in ip_list}

        for future in as_completed(future_to_ip):
            ip = future_to_ip[future]
            try:
                result = future.result()
            except Exception as e:
                result = {"success": False, "ip": ip, "error": str(e)}
            import time; time.sleep(2)  # 2s delay between switches to avoid RADIUS lockout

            results.append(result)
            done += 1

            if progress_callback:
                progress_callback(done, total, ip, result)

    return results