"""
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}")


# ── Multi-shelf IP collection ─────────────────────────────────────────────────
def _collect_shelves(iom_choice: str) -> list:
    """
    Prompt for the password and IOM IP addresses one shelf at a time,
    offering to add more shelves after each entry.

    Each shelf has its own password because the Admin password is the BMC
    serial number, which is unique to each physical shelf.

    Returns a list of (password, [(iom, ip), ...]) tuples, one per shelf.
    """
    shelves   = []
    shelf_num = 1

    while True:
        info(f"Shelf {shelf_num} — enter password and IP address(es).")
        password = prompt_password()
        shelf = []
        if iom_choice in ("1", "3"):
            ip = prompt_ip(f"  Shelf {shelf_num} IOM1 IP address")
            shelf.append(("IOM1", ip))
        if iom_choice in ("2", "3"):
            ip = prompt_ip(f"  Shelf {shelf_num} IOM2 IP address")
            shelf.append(("IOM2", ip))
        shelves.append((password, shelf))
        print()

        if not prompt_yn("Add another shelf?", default=False):
            break
        shelf_num += 1
        print()

    return shelves


def _make_targets(shelves: list) -> list:
    """
    Convert the shelves structure into a flat list of (label, iom, ip, password)
    tuples suitable for _show_fw_versions().  When there is only one shelf the
    label is just the IOM name; for multiple shelves it includes the shelf number.
    """
    multi = len(shelves) > 1
    return [
        (f"S{i} / {iom}" if multi else iom, iom, ip, password)
        for i, (password, shelf) in enumerate(shelves, 1)
        for iom, ip in shelf
    ]


# ── 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.
    Supports updating multiple shelves in a single run.
    """
    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()

    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()

    # ── Collect IPs for one or more shelves ───────────────────────────────────
    shelves = _collect_shelves(iom_choice)
    targets = _make_targets(shelves)

    rule("Current Firmware Versions")
    _show_fw_versions(targets)

    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(shelves) > 1:
        warn(f"Updating {len(shelves)} shelves sequentially — Shelf 1 first.")
    elif any(len(shelf) > 1 for shelf in shelves):
        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

    # ── Run updates shelf by shelf, IOM by IOM ────────────────────────────────
    multi_shelf = len(shelves) > 1
    for i, (password, shelf) in enumerate(shelves, 1):
        for iom, ip in shelf:
            heading = f"Shelf {i} — {iom} ({ip})" if multi_shelf else f"{iom} ({ip})"
            rule(heading)
            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(targets)

    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()