#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
列出所有端口使用情况（跨平台）。

优先使用 psutil（若安装），否则基于系统自带工具：
- Windows: netstat + tasklist + wmic
- Linux: ss（优先）/netstat + ps
- macOS: lsof + ps
"""

from __future__ import annotations

import platform
import re
import shutil
import subprocess
from typing import List, Tuple, Optional, Dict


def _run_command(command_args: List[str]) -> Tuple[int, str]:
    """运行命令并返回 (returncode, stdout)"""
    process = subprocess.Popen(
        command_args,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding='utf-8',
        errors='replace',
    )
    assert process.stdout is not None
    output_lines = []
    for line in process.stdout:
        output_lines.append(line)
    process.wait()
    return process.returncode, ''.join(output_lines)


def _get_all_process_info_windows() -> Dict[int, Tuple[str, str]]:
    """批量获取所有 Windows 进程信息（进程名和命令行）
    
    返回: Dict[pid, (name, cmdline)]
    """
    pid_info: Dict[int, Tuple[str, str]] = {}
    
    # 方法1: 尝试使用 wmic 批量获取（最快）
    try:
        rc, out = _run_command(["wmic", "process", "get", "ProcessId,Name,CommandLine", "/format:csv"])
        if rc == 0:
            lines = out.splitlines()
            for line in lines[1:]:  # 跳过表头
                if not line.strip():
                    continue
                # CSV 格式: Node,ProcessId,Name,CommandLine
                parts = line.split(',')
                if len(parts) >= 4:
                    try:
                        pid = int(parts[1])
                        name = parts[2].strip()
                        cmdline = ','.join(parts[3:]).strip()  # CommandLine 可能包含逗号
                        if pid > 0:
                            pid_info[pid] = (name, cmdline)
                    except (ValueError, IndexError):
                        continue
            if pid_info:
                return pid_info
    except Exception:
        pass
    
    # 方法2: 使用 PowerShell 批量获取（备用方案）
    try:
        ps_cmd = "Get-WmiObject Win32_Process | Select-Object ProcessId,Name,CommandLine | ConvertTo-Csv -NoTypeInformation"
        rc, out = _run_command(["powershell", "-Command", ps_cmd])
        if rc == 0:
            lines = out.splitlines()
            for line in lines[1:]:  # 跳过表头
                if not line.strip():
                    continue
                # CSV 格式: "ProcessId","Name","CommandLine"
                # 使用正则表达式解析带引号的 CSV
                match = re.match(r'"(\d+)","([^"]*)","(.*)"', line)
                if match:
                    try:
                        pid = int(match.group(1))
                        name = match.group(2)
                        cmdline = match.group(3)
                        if pid > 0:
                            pid_info[pid] = (name, cmdline)
                    except (ValueError, IndexError):
                        continue
            if pid_info:
                return pid_info
    except Exception:
        pass
    
    # 方法3: 使用 tasklist 获取进程名（至少能获取进程名）
    try:
        rc, out = _run_command(["tasklist", "/FO", "CSV", "/NH"])
        if rc == 0:
            for line in out.splitlines():
                if not line.strip():
                    continue
                # CSV 格式: "Image Name","PID","Session Name","Session#","Mem Usage"
                parts = line.split('","')
                if len(parts) >= 2:
                    try:
                        name = parts[0].strip('"')
                        pid = int(parts[1].strip('"'))
                        if pid > 0:
                            # 如果没有命令行信息，至少保存进程名
                            if pid not in pid_info:
                                pid_info[pid] = (name, "")
                            else:
                                # 更新进程名
                                old_name, old_cmdline = pid_info[pid]
                                pid_info[pid] = (name, old_cmdline)
                    except (ValueError, IndexError):
                        continue
    except Exception:
        pass
    
    return pid_info


def _get_all_process_info_linux(pids: set) -> Dict[int, str]:
    """批量获取 Linux 进程的命令行
    
    返回: Dict[pid, cmdline]
    """
    pid_cmdline: Dict[int, str] = {}
    for pid in pids:
        try:
            # 读取 /proc/<pid>/cmdline
            with open(f"/proc/{pid}/cmdline", "r", encoding='utf-8', errors='replace') as f:
                cmdline = f.read().replace('\0', ' ').strip()
                if cmdline:
                    pid_cmdline[pid] = cmdline
        except Exception:
            pass
    return pid_cmdline


def _get_all_process_info_mac() -> Dict[int, str]:
    """批量获取所有 macOS 进程的命令行
    
    返回: Dict[pid, cmdline]
    """
    pid_cmdline: Dict[int, str] = {}
    try:
        # 一次性获取所有进程的命令行
        rc, out = _run_command(["ps", "-eo", "pid,command"])
        if rc == 0:
            for line in out.splitlines():
                if not line.strip():
                    continue
                # 格式: PID COMMAND
                parts = line.split(None, 1)  # 只分割第一个空格
                if len(parts) >= 2:
                    try:
                        pid = int(parts[0])
                        cmdline = parts[1].strip()
                        if pid > 0:
                            pid_cmdline[pid] = cmdline
                    except (ValueError, IndexError):
                        continue
    except Exception:
        pass
    return pid_cmdline


def _try_psutil() -> List[Tuple[int, int, str, str, str]]:
    """使用 psutil 获取所有端口使用情况
    
    返回: List[(port, pid, name, status, cmdline)]
    """
    try:
        import psutil  # type: ignore
    except Exception:
        return []

    results: List[Tuple[int, int, str, str, str]] = []
    seen = set()  # (port, pid) 去重
    
    try:
        for conn in psutil.net_connections(kind='inet'):
            laddr = conn.laddr if conn.laddr else None
            if not laddr:
                continue
            
            port = getattr(laddr, 'port', None)
            if port is None:
                continue
            
            pid = conn.pid
            if pid is None:
                continue
            
            # 去重
            key = (port, pid)
            if key in seen:
                continue
            seen.add(key)
            
            # 获取进程信息
            try:
                proc = psutil.Process(pid)
                name = proc.name()
                try:
                    cmdline = ' '.join(proc.cmdline())
                except Exception:
                    cmdline = ""
                status = conn.status if hasattr(conn, 'status') else ""
            except Exception:
                name = "<unknown>"
                cmdline = ""
                status = ""
            
            results.append((port, pid, name, status, cmdline))
    except Exception:
        return []
    
    return results


def _windows_list() -> List[Tuple[int, int, str, str, str]]:
    """Windows 系统列出所有端口使用情况
    
    返回: List[(port, pid, name, status, cmdline)]
    """
    results: List[Tuple[int, int, str, str, str]] = []
    seen = set()  # (port, pid) 去重
    
    # 获取所有端口连接
    rc, out = _run_command(["netstat", "-ano"])
    if rc != 0:
        return results
    
    # 解析 netstat 输出
    # 格式: TCP    0.0.0.0:8080    0.0.0.0:0    LISTENING    1234
    pattern = re.compile(r"^\s*(TCP|UDP)\s+[^\s]*:(\d+)\s+.*\s+(\d+)\s*$", re.IGNORECASE)
    
    port_pid_map: Dict[Tuple[int, int], str] = {}  # {(port, pid): status}
    
    for line in out.splitlines():
        m = pattern.match(line)
        if not m:
            continue
        protocol = m.group(1).upper()
        port = int(m.group(2))
        pid = int(m.group(3))
        
        # 提取状态（如果有）
        status = ""
        if "LISTENING" in line:
            status = "LISTEN"
        elif "ESTABLISHED" in line:
            status = "ESTABLISHED"
        elif "TIME_WAIT" in line:
            status = "TIME_WAIT"
        
        key = (port, pid)
        if key not in port_pid_map:
            port_pid_map[key] = status
    
    # 批量获取所有进程信息（优化：一次性获取，而不是逐个查询）
    all_pids = {pid for (port, pid), status in port_pid_map.items()}
    pid_info = _get_all_process_info_windows()
    
    # 构建结果
    for (port, pid), status in port_pid_map.items():
        if pid in pid_info:
            name, cmdline = pid_info[pid]
        else:
            name = "<unknown>"
            cmdline = ""
        results.append((port, pid, name, status, cmdline))
    
    return results


def _linux_list() -> List[Tuple[int, int, str, str, str]]:
    """Linux 系统列出所有端口使用情况
    
    返回: List[(port, pid, name, status, cmdline)]
    """
    results: List[Tuple[int, int, str, str, str]] = []
    seen = set()  # (port, pid) 去重
    
    # 优先使用 ss
    if shutil.which("ss") is not None:
        rc, out_tcp = _run_command(["ss", "-ltnp"])  # TCP
        rc2, out_udp = _run_command(["ss", "-lunp"])  # UDP
        combined = (out_tcp if rc == 0 else "") + "\n" + (out_udp if rc2 == 0 else "")
        
        # 解析 ss 输出
        # 示例: LISTEN 0 128 0.0.0.0:8080 0.0.0.0:* users:(("nginx",pid=1234,fd=6))
        for line in combined.splitlines():
            if not line.strip() or "State" in line:
                continue
            
            # 提取端口
            port_match = re.search(r":(\d+)(?:\s|$)", line)
            if not port_match:
                continue
            port = int(port_match.group(1))
            
            # 提取状态
            status = ""
            if "LISTEN" in line:
                status = "LISTEN"
            elif "ESTAB" in line:
                status = "ESTABLISHED"
            
            # 提取进程信息
            proc_match = re.search(r'users:\(\(\"([^\"]+)\",pid=(\d+)', line)
            if proc_match:
                name = proc_match.group(1)
                pid = int(proc_match.group(2))
                
                key = (port, pid)
                if key in seen:
                    continue
                seen.add(key)
                
                # 先收集所有 PID，稍后批量获取命令行
                results.append((port, pid, name, status, ""))  # 临时空命令行
        
        # 批量获取所有进程的命令行（优化：一次性获取）
        if results:
            all_pids = {pid for port, pid, name, status, cmdline in results}
            pid_cmdline_map = _get_all_process_info_linux(all_pids)
            
            # 更新结果中的命令行
            final_results = []
            for port, pid, name, status, _ in results:
                cmdline = pid_cmdline_map.get(pid, "")
                final_results.append((port, pid, name, status, cmdline))
            return final_results
    
    # 回退使用 netstat
    if shutil.which("netstat") is not None:
        rc, out = _run_command(["netstat", "-tunlp"])
        if rc == 0:
            for line in out.splitlines():
                if not line.strip() or "Proto" in line:
                    continue
                
                # 提取端口和进程
                # 示例: tcp  0  0  0.0.0.0:8080  0.0.0.0:*  LISTEN  1234/nginx
                port_match = re.search(r":(\d+)(?:\s|$)", line)
                if not port_match:
                    continue
                port = int(port_match.group(1))
                
                proc_match = re.search(r"\s(\d+)/(\S+)\s*$", line)
                if not proc_match:
                    continue
                
                pid = int(proc_match.group(1))
                name = proc_match.group(2)
                
                key = (port, pid)
                if key in seen:
                    continue
                seen.add(key)
                
                # 提取状态
                status = ""
                if "LISTEN" in line:
                    status = "LISTEN"
                elif "ESTABLISHED" in line:
                    status = "ESTABLISHED"
                
                # 先收集所有 PID，稍后批量获取命令行
                results.append((port, pid, name, status, ""))  # 临时空命令行
    
    # 批量获取所有进程的命令行（优化：一次性获取）
    if results:
        all_pids = {pid for port, pid, name, status, cmdline in results}
        pid_cmdline_map = _get_all_process_info_linux(all_pids)
        
        # 更新结果中的命令行
        final_results = []
        for port, pid, name, status, _ in results:
            cmdline = pid_cmdline_map.get(pid, "")
            final_results.append((port, pid, name, status, cmdline))
        return final_results
    
    return results


def _mac_list() -> List[Tuple[int, int, str, str, str]]:
    """macOS 系统列出所有端口使用情况
    
    返回: List[(port, pid, name, status, cmdline)]
    """
    results: List[Tuple[int, int, str, str, str]] = []
    seen = set()  # (port, pid) 去重
    
    if shutil.which("lsof") is None:
        return results
    
    rc, out = _run_command(["lsof", "-nP", "-iTCP", "-iUDP"])
    if rc != 0:
        return results
    
    # 解析 lsof 输出
    # 格式: COMMAND PID USER FD TYPE DEVICE SIZE/OFF NODE NAME
    # 示例: nginx 1234 user 6u IPv4 0x... TCP *:8080 (LISTEN)
    lines = [l for l in out.splitlines() if l.strip()]
    if not lines:
        return results
    
    for line in lines[1:]:  # 跳过表头
        parts = re.split(r"\s+", line)
        if len(parts) < 9:
            continue
        
        name = parts[0]
        try:
            pid = int(parts[1])
        except ValueError:
            continue
        
        # 提取端口和状态
        name_part = ' '.join(parts[8:])
        port_match = re.search(r":(\d+)", name_part)
        if not port_match:
            continue
        port = int(port_match.group(1))
        
        status = ""
        if "LISTEN" in name_part:
            status = "LISTEN"
        elif "ESTABLISHED" in name_part:
            status = "ESTABLISHED"
        
        key = (port, pid)
        if key in seen:
            continue
        seen.add(key)
        
        # 先收集所有 PID，稍后批量获取命令行
        results.append((port, pid, name, status, ""))  # 临时空命令行
    
    # 批量获取所有进程的命令行（优化：一次性获取）
    if results:
        all_pids = {pid for port, pid, name, status, cmdline in results}
        pid_cmdline_map = _get_all_process_info_mac()
        
        # 更新结果中的命令行
        final_results = []
        for port, pid, name, status, _ in results:
            cmdline = pid_cmdline_map.get(pid, "")
            final_results.append((port, pid, name, status, cmdline))
        return final_results
    
    return results


def get_all_ports() -> List[Tuple[int, int, str, str, str]]:
    """获取所有端口使用情况（供 whoport 等模块复用）
    
    返回: List[(port, pid, name, status, cmdline)]
    """
    results = _try_psutil()
    if not results:
        system = platform.system().lower()
        if system == 'windows':
            results = _windows_list()
        elif system == 'linux':
            results = _linux_list()
        elif system == 'darwin':
            results = _mac_list()
        else:
            results = []
    results.sort(key=lambda x: x[0])
    return results


def command_port_list(full: bool = False) -> None:
    """列出所有端口使用情况
    
    Args:
        full: 如果为 True，显示完整的命令行（不截断）
    """
    import sys
    
    results = get_all_ports()
    
    if not results:
        print("未发现端口使用情况或缺少权限/工具无法识别。")
        print("建议：以管理员/root 权限运行，或安装 psutil。")
        return
    
    # 处理 Windows 控制台编码问题 - 使用安全的输出函数
    def safe_print(text: str) -> None:
        """安全打印，处理编码问题"""
        try:
            print(text)
        except UnicodeEncodeError:
            # 如果遇到编码问题，使用 errors='replace' 重新编码
            try:
                if sys.stdout.encoding:
                    encoded = text.encode(sys.stdout.encoding, errors='replace')
                    sys.stdout.buffer.write(encoded)
                    sys.stdout.buffer.write(b'\n')
                else:
                    # 备用方案：使用 ASCII
                    safe_text = text.encode('ascii', 'replace').decode('ascii')
                    print(safe_text)
            except Exception:
                # 最后的备用方案
                print(text.encode('ascii', 'replace').decode('ascii', 'replace'))
    
    # 打印表头
    try:
        safe_print(f"{'端口':<8} {'PID':<8} {'进程名':<20} {'状态':<12} {'命令行'}")
    except Exception:
        safe_print(f"{'Port':<8} {'PID':<8} {'Process':<20} {'Status':<12} {'Command'}")
    safe_print("-" * 150)
    
    # 打印结果
    for port, pid, name, status, cmdline in results:
        # 根据 full 参数决定是否截断命令行
        if not full:
            # 默认截断到150字符，足够显示大部分完整路径
            max_cmdline_len = 150
            if len(cmdline) > max_cmdline_len:
                cmdline = cmdline[:max_cmdline_len - 3] + "..."
        # 如果 full=True，显示完整命令行，不截断
        status_str = status if status else "-"
        safe_print(f"{port:<8} {pid:<8} {name:<20} {status_str:<12} {cmdline}")

