#!/bin/bash

# 配置项
VM_XML_SOURCE_PATH="/etc/vmdeploy/Windows_11.xml"
SERIAL_PASSTHROUGH_XML="/tmp/serial_passthrough.xml"
LOG_FILE="/tmp/serial_direct.log"

# 定义全局数组
SERIAL_PARAMS=()

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# 日志函数
log() {
    echo "[$(date +%Y-%m-%d_%H:%M:%S)] $1" >> "$LOG_FILE"
}

# 已验证资源，命令中COM口仅做直通标识，无实际意义
convert_com_to_params() {
    local com_num=$1
    case "$com_num" in
        COM1)
            SERIAL_PARAMS=("0" "0x3E8" "4")  # target_port=0, iobase=0x3E8, irq=4
            ;;
        COM2)
            SERIAL_PARAMS=("1" "0x2E0" "7")  # target_port=1, iobase=0x2E0, irq=7
            ;;
        COM3)
            SERIAL_PARAMS=("2" "0x2E8" "3")  # target_port=2, iobase=0x2E8, irq=3
            ;;
        *)
            echo -e "${RED}Error: Only COM1/COM2/COM3 are supported (max 3 serial devices)${NC}"
            log "Unsupported COM port: $com_num (only COM1/2/3 are allowed)"
            return 1
            ;;
    esac
    return 0
}

# 校验参数合法性
validate_params() {
    # 1. 校验参数数量
    if [ $# -ne 2 ]; then
        echo -e "${RED}Error: Need exactly 2 parameters (serial device path + COM port), e.g., /dev/ttyS2 COM1${NC}"
        log "Invalid parameter count: $# (expected 2)"
        return 1
    fi

    local serial_dev=$1
    local com_port=$2

    # 2. 校验串口号
    if ! echo "$com_port" | grep -qE '^COM[123]$'; then
        echo -e "${RED}Error: COM port must be COM1/COM2/COM3 (max 3 serial devices)${NC}"
        log "Invalid COM port: $com_port (only COM1/2/3 are allowed)"
        return 1
    fi

    # 3. 校验宿主机串口设备是否存在
    if [ ! -c "$serial_dev" ]; then
        echo -e "${RED}Error: Serial device $serial_dev does not exist (check /dev/ttyS* list)${NC}"
        log "Serial device not found: $serial_dev"
        return 1
    fi

    return 0
}

# 统计当前已直通的串口数量
count_existing_serials() {
    # 提取所有<serial type='dev'>块并统计数量
    local serial_count=$(sed -n '/<serial type='\''dev'\''>/,/<\/serial>/p' "$VM_XML_SOURCE_PATH" | grep -c "<serial type='dev'>")
    echo "$serial_count"
}

# 冲突检查：1、设备是否已直通 2、iobase或irq任意一项重复即判定冲突
check_serial_existed() {
    local serial_dev=$1
    local iobase=$2
    local com_port=$3

    # 检查0：最多允许直通3路串口，超过则直接判定冲突
    local existing_serial_count=$(count_existing_serials)
    if [ "$existing_serial_count" -ge 3 ]; then
        echo -e "${RED}Error: Max 3 serial devices are allowed to pass through (current: $existing_serial_count)${NC}"
        log "Max serial passthrough limit reached (3), current count: $existing_serial_count"
        return 0
    fi

    # 检查1：设备路径是否已在XML中
    # 步骤：1. 提取所有<serial type='dev'>块 2. 检查是否包含该设备路径
    local dev_existed=$(sed -n '/<serial type='\''dev'\''>/,/<\/serial>/p' "$VM_XML_SOURCE_PATH" | grep -c "<source path='$serial_dev'/>")
    if [ $dev_existed -gt 0 ]; then
        echo -e "${RED}Error: Serial device $serial_dev is already passed through to another COM port!${NC}"
        log "Serial device $serial_dev already exists in XML (bound to other COM port), skip duplicate config"
        return 0
    fi

    # 检查2：iobase是否被占用（避免地址冲突）
    local iobase_existed=$(grep -c "<address type='isa' iobase='$iobase' " "$VM_XML_SOURCE_PATH")
    if [ $iobase_existed -gt 0 ]; then
        echo -e "${RED}Error: iobase $iobase (for $com_port) is already occupied${NC}"
        log "iobase $iobase is occupied, cannot pass through $serial_dev to $com_port"
        return 0
    fi

    # 检查3：irq是否被占用
    local irq_existed=$(grep -c "<address type='isa' .* irq='$irq' " "$VM_XML_SOURCE_PATH")
    if [ $irq_existed -gt 0 ]; then
        echo -e "${RED}Error: irq $irq (for $com_port) is already occupied${NC}"
        log "irq $irq is occupied, cannot pass through $serial_dev to $com_port"
        return 0
    fi

    return 1
}

# 生成串口直通XML片段
generate_serial_xml() {
    local serial_dev=$1
    local com_port=$2
    local target_port=$3
    local iobase=$4
    local irq=$5

    # 生成标准化串口直通XML
    cat > "$SERIAL_PASSTHROUGH_XML" << EOF
<!-- Serial passthrough config: $serial_dev → $com_port -->
<serial type='dev'>
  <source path='$serial_dev'/>
  <target type='isa-serial' port='$target_port'>
    <model name='isa-serial'/>
  </target>
  <address type='isa' iobase='$iobase' irq='$irq'/>
</serial>
EOF
    log "Generated serial passthrough XML: $serial_dev → $com_port (target_port=$target_port, iobase=$iobase, irq=$irq)"
}

# 移除原有串口配置
remove_existing_serial_config() {
    local serial_dev=$1
    sed -i "#<!-- Serial passthrough config: $serial_dev#,#<\/serial>#d" "$VM_XML_SOURCE_PATH"
    log "Removed existing serial config for $serial_dev (if any)"
}

# 插入串口XML到memballoon上方
insert_serial_xml() {
    local serial_dev=$1
    local com_port=$2
    local target_port=$3
    local iobase=$4
    local irq=$5

    # 1. 备份原XML
    local backup_xml="${VM_XML_SOURCE_PATH}.bak_$(date +%Y%m%d_%H%M%S)"
    cp "$VM_XML_SOURCE_PATH" "$backup_xml"
    log "Backed up VM XML to: $backup_xml"
    echo -e "${YELLOW}Backed up original VM XML to: $backup_xml${NC}"

    # 2. 移除旧配置
    remove_existing_serial_config "$serial_dev"

    # 3. 生成串口XML
    generate_serial_xml "$serial_dev" "$com_port" "$target_port" "$iobase" "$irq"

    # 4. 插入到<memballoon>上方
    grep -v '^$' "$SERIAL_PASSTHROUGH_XML" | sed -r 's/\x1B\[[0-9;]*m//g' > /tmp/tmp_serial_xml
    local line_num=$(grep -n "<memballoon model='virtio'>" "$VM_XML_SOURCE_PATH" | cut -d: -f1)
    if [ -n "$line_num" ]; then
        local insert_line=$((line_num - 1))  # 插入到memballoon行之前
        sed -i "${insert_line}r /tmp/tmp_serial_xml" "$VM_XML_SOURCE_PATH"
    else
        echo -e "${RED}Error: <memballoon model='virtio'> not found in XML${NC}"
        log "memballoon node not found, insert failed"
        cp "$backup_xml" "$VM_XML_SOURCE_PATH"
        return 1
    fi

    # 5. 清理临时文件
    rm -f /tmp/tmp_serial_xml

    # 6. 验证插入结果
    if grep -q "<!-- Serial passthrough config: $serial_dev → $com_port -->" "$VM_XML_SOURCE_PATH"; then
        echo -e "${GREEN}Successfully inserted serial passthrough config: $serial_dev → $com_port${NC}"
        log "Serial config inserted successfully: $serial_dev → $com_port"
        return 0
    else
        echo -e "${RED}Failed to insert serial config, restoring backup${NC}"
        cp "$backup_xml" "$VM_XML_SOURCE_PATH"
        log "Insert failed, restored backup: $backup_xml"
        return 1
    fi
}

# 主逻辑
main() {
    # 初始化日志
    > "$LOG_FILE"
    log "===== Serial passthrough script started ====="

    # 1. 参数校验
    if ! validate_params "$@"; then
        exit 1
    fi

    local serial_dev=$1
    local com_port=$2
    echo -e "${YELLOW}Processing serial device: $serial_dev → $com_port${NC}"
    log "Start processing serial device: $serial_dev → $com_port"

    # 2. 转换COM口到硬件参数
    if ! convert_com_to_params "$com_port"; then
        exit 1
    fi
    local target_port="${SERIAL_PARAMS[0]}"
    local iobase="${SERIAL_PARAMS[1]}"
    local irq="${SERIAL_PARAMS[2]}"
    echo -e "${YELLOW}Mapped COM params: target_port=$target_port, iobase=$iobase, irq=$irq${NC}"
    log "Mapped COM params: $com_port → target_port=$target_port, iobase=$iobase, irq=$irq"

    # 3. 冲突检查
    if check_serial_existed "$serial_dev" "$iobase" "$com_port"; then
        echo -e "${YELLOW}Skipping serial device $serial_dev (conflict detected)${NC}"
        log "Skipping serial device $serial_dev due to conflict"
        exit 0
    fi

    # 4. 插入XML配置
    if ! insert_serial_xml "$serial_dev" "$com_port" "$target_port" "$iobase" "$irq"; then
        echo -e "${RED}Failed to process serial device $serial_dev${NC}"
        log "Failed to process serial device $serial_dev"
        exit 1
    fi

    # 5. 清理临时文件
    rm -f "$SERIAL_PASSTHROUGH_XML"
    log "===== Serial passthrough script completed successfully ====="
    echo -e "${GREEN}=====================${NC}"
    echo -e "${GREEN}Serial passthrough config is in effect${NC}"
    echo -e "${GREEN}Device: $serial_dev → Windows COM port: $com_port${NC}"
    echo -e "${GREEN}=====================${NC}"
    exit 0
}

# 执行主逻辑（接收vmdeploy传递的2个参数：设备路径 + 串口号）
main "$@"

