"""
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 IOM IP addresses one shelf at a time, offering to add more
    shelves after each entry.

    Returns a list of shelves; each shelf is a list of (iom, ip) tuples
    for the IOM(s) selected (e.g. [("IOM1", "10.0.0.1")] or
    [("IOM1", "10.0.0.1"), ("IOM2", "10.0.0.2")]).
    """
    shelves   = []
    shelf_num = 1

    while True:
        info(f"Enter the management IP address(es) for Shelf {shelf_num}.")
        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(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) 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)
        for i, 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()

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

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

    rule("Current Firmware Versions")
    _show_fw_versions(password, 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, 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(password, 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()