"""
workflow_firmware.py — IOM and Fabric Card firmware update workflow.
Connects to each IOM via its network IP address using the Redfish API.
"""

import os
import time

from redfish import (
    _redfish_upload_firmware,
    _redfish_trigger_update,
    _redfish_poll_tasks,
    _redfish_restart_iom,
    _redfish_reset_fabric,
    _show_fw_versions,
)
from ui import (
    _c, C,
    banner, rule, draw_box,
    info, ok, warn, error,
    prompt, prompt_ip, prompt_yn, prompt_password,
)


# ── Firmware file selection helper ────────────────────────────────────────────
def _prompt_fw_file(label: str) -> str:
    """
    Scan the current working directory for firmware files and let the user
    pick one by number, or enter a custom path as the last option.
    Files are sorted most-recently-modified first.
    """
    cwd = os.getcwd()
    FW_EXTS = {".bin", ".img", ".fw", ".hex", ".zip", ".tar", ".tgz", ".gz"}

    try:
        candidates = sorted(
            [f for f in os.listdir(cwd)
             if not f.startswith(".")
             and os.path.isfile(os.path.join(cwd, f))
             and os.path.splitext(f)[1].lower() in FW_EXTS],
            key=lambda f: os.path.getmtime(os.path.join(cwd, f)),
            reverse=True,
        )
    except OSError:
        candidates = []

    print()
    if candidates:
        info(f"Firmware files found in {cwd}:")
        for i, fname in enumerate(candidates, 1):
            sz = os.path.getsize(os.path.join(cwd, fname))
            print(f"    {_c(C.BOLD, str(i))}  {fname}  {_c(C.DIM, f'({sz // 1024} KB)')}")
        custom_idx = len(candidates) + 1
        print(f"    {_c(C.BOLD, str(custom_idx))}  Enter a custom file path")
        print()

        while True:
            choice = prompt(f"Select {label} [1-{custom_idx}]")
            if choice.isdigit():
                idx = int(choice)
                if 1 <= idx <= len(candidates):
                    path = os.path.join(cwd, candidates[idx - 1])
                    sz   = os.path.getsize(path)
                    ok(f"Selected: {candidates[idx - 1]} ({sz // 1024} KB)")
                    return path
                if idx == custom_idx:
                    break
            warn(f"Please enter a number between 1 and {custom_idx}.")
    else:
        info(f"No firmware files found in {cwd}.")

    # Manual path entry
    while True:
        path = prompt(f"Path to {label}")
        if os.path.isfile(path):
            sz = os.path.getsize(path)
            ok(f"File: {path} ({sz // 1024} KB)")
            return path
        warn(f"File not found: {path}")


# ── Per-IOM update helpers ────────────────────────────────────────────────────
def _update_iom_fw(password: str, ip: str, iom: str, fw_path: str) -> bool:
    """Upload and apply IOM firmware for one IOM, then restart it."""
    sz = os.path.getsize(fw_path)
    info(f"Uploading IOM firmware ({sz // 1024} KB) to {iom} at {ip}...")
    ok_flag, data = _redfish_upload_firmware(password, ip, fw_path)
    if not ok_flag:
        error(f"Upload failed: {data}")
        return False
    ok("Firmware file uploaded.")

    info(f"Triggering {iom} firmware update...")
    ok_flag, data = _redfish_trigger_update(
        password, ip, f"/redfish/v1/Managers/{iom}",
    )
    if not ok_flag:
        error(f"Update trigger failed: {data}")
        return False
    ok("Update triggered.")

    info("Monitoring update progress (this may take several minutes)...")
    ok_flag, msg = _redfish_poll_tasks(password, ip)
    if not ok_flag:
        warn(f"Task monitoring ended: {msg}")
    else:
        ok(msg)

    info(f"Restarting {iom}...")
    _redfish_restart_iom(password, ip, iom)   # connection drop on restart is normal
    ok(f"{iom} restart initiated. Waiting 30s for reboot...")
    time.sleep(30)
    return True


def _update_fabric_fw(password: str, ip: str, iom: str, fw_path: str) -> bool:
    """
    Upload and apply Fabric Card firmware for one IOM.
    Per the service guide, the firmware file must be re-uploaded even if it was
    already uploaded during the IOM firmware step.
    After the update: restart fabric card, then restart IOM.
    """
    sz = os.path.getsize(fw_path)
    info(f"Uploading Fabric Card firmware ({sz // 1024} KB) to {iom} at {ip}...")
    ok_flag, data = _redfish_upload_firmware(password, ip, fw_path)
    if not ok_flag:
        error(f"Upload failed: {data}")
        return False
    ok("Firmware file uploaded.")

    info(f"Triggering {iom} Fabric Card firmware update...")
    ok_flag, data = _redfish_trigger_update(
        password, ip, f"/redfish/v1/Chassis/{iom}/NetworkAdapters/1",
    )
    if not ok_flag:
        error(f"Update trigger failed: {data}")
        return False
    ok("Update triggered.")

    info("Monitoring update progress...")
    ok_flag, msg = _redfish_poll_tasks(password, ip)
    if not ok_flag:
        warn(f"Task monitoring ended: {msg}")
    else:
        ok(msg)

    info(f"Restarting {iom} Fabric Card...")
    _redfish_reset_fabric(password, ip, iom)
    ok("Fabric Card restart initiated. Waiting 15s...")
    time.sleep(15)

    info(f"Restarting {iom} after Fabric Card update...")
    _redfish_restart_iom(password, ip, iom)
    ok(f"{iom} restart initiated. Waiting 30s for reboot...")
    time.sleep(30)
    return True


# ── Firmware Update Workflow ──────────────────────────────────────────────────
def firmware_update_workflow():
    """
    Standalone firmware update for IOM and Fabric Card firmware.
    Connects to each IOM via its network IP (not serial loopback) — uploading
    firmware over 115200-baud serial would be impractically slow.
    """
    banner()
    rule("IOM & Fabric Card Firmware Update")
    info("This procedure connects to each IOM via its network IP address.")
    info("Ensure this system has network access to the IOM management interface.")
    print()

    password = prompt_password()
    print()

    print("  Which IOM(s) would you like to update?")
    print(f"    {_c(C.BOLD, '1')}  IOM1 only")
    print(f"    {_c(C.BOLD, '2')}  IOM2 only")
    print(f"    {_c(C.BOLD, '3')}  Both IOM1 and IOM2")
    print()

    while True:
        iom_choice = prompt("Select option [1-3]")
        if iom_choice in ("1", "2", "3"):
            break
        warn("Please enter 1, 2, or 3.")
    print()

    info("Enter the management IP address for each IOM to update.")
    iom1_ip = prompt_ip("  IOM1 IP address") if iom_choice in ("1", "3") else ""
    iom2_ip = prompt_ip("  IOM2 IP address") if iom_choice in ("2", "3") else ""
    print()

    ioms = []
    if iom_choice in ("1", "3"):
        ioms.append(("IOM1", iom1_ip))
    if iom_choice in ("2", "3"):
        ioms.append(("IOM2", iom2_ip))

    rule("Current Firmware Versions")
    _show_fw_versions(password, ioms)

    print("  What would you like to update?")
    print(f"    {_c(C.BOLD, '1')}  IOM Firmware only")
    print(f"    {_c(C.BOLD, '2')}  Fabric Card Firmware only")
    print(f"    {_c(C.BOLD, '3')}  Both IOM and Fabric Card Firmware")
    print(f"    {_c(C.BOLD, '4')}  Cancel")
    print()

    while True:
        choice = prompt("Select option [1-4]")
        if choice in ("1", "2", "3", "4"):
            break
        warn("Please enter 1, 2, 3, or 4.")

    if choice == "4":
        info("Firmware update cancelled.")
        return

    update_iom    = choice in ("1", "3")
    update_fabric = choice in ("2", "3")
    iom_fw_path    = ""
    fabric_fw_path = ""

    if update_iom:
        iom_fw_path = _prompt_fw_file("IOM firmware file")

    if update_fabric:
        fabric_fw_path = _prompt_fw_file("Fabric Card firmware file")

    print()
    warn("For HA systems: update the passive IOM first.")
    if len(ioms) > 1:
        warn("IOM1 will be updated first — adjust order if IOM2 is passive.")
    print()

    if not prompt_yn("Proceed with firmware update?", default=True):
        info("Firmware update cancelled.")
        return

    for iom, ip in ioms:
        rule(f"{iom} ({ip})")
        if update_iom:
            _update_iom_fw(password, ip, iom, iom_fw_path)
        if update_fabric:
            _update_fabric_fw(password, ip, iom, fabric_fw_path)

    rule("Post-Update Firmware Validation")
    _show_fw_versions(password, ioms)

    print()
    draw_box([
        f"{_c(C.YEL, 'IMPORTANT -- For HA (Dual-Controller) Systems:')}",
        "",
        "After updating this controller's IOMs:",
        "  1. Log into TrueNAS and initiate a failover.",
        "  2. Re-run this tool to update the other controller.",
    ], colour=C.YEL)
    print()