"""
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,
    _wait_for_iom_online,
    _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 firmware/ directory (next to es24n_conf.py) 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.
    """
    fw_dir = os.path.normpath(
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "firmware")
    )
    FW_EXTS = {".bin", ".img", ".fw", ".fwc", ".hex", ".zip", ".tar", ".tgz", ".gz"}

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

    print()
    if candidates:
        info(f"Firmware files found in {fw_dir}:")
        for i, fname in enumerate(candidates, 1):
            sz = os.path.getsize(os.path.join(fw_dir, 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(fw_dir, 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 {fw_dir}.")

    # 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() -> list:
    """
    Prompt for the password and a single IOM IP address per shelf,
    offering to add more shelves after each entry.

    Either IOM's IP can be used — both IOMs share the same Redfish endpoint
    and can reach each other's resources.  Each shelf has its own password
    because the Admin password is the BMC serial number, unique per shelf.

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

    while True:
        info(f"Shelf {shelf_num} — enter password and IOM IP address.")
        password = prompt_password()
        ip = prompt_ip(f"  Shelf {shelf_num} IOM IP address (IOM1 or IOM2)")
        shelves.append((password, ip))
        print()

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

    return shelves


def _make_targets(shelves: list, iom_choice: str) -> 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.
    All IOMs in a shelf share the same IP connection.
    """
    ioms = []
    if iom_choice in ("1", "3"):
        ioms.append("IOM1")
    if iom_choice in ("2", "3"):
        ioms.append("IOM2")

    multi = len(shelves) > 1
    return [
        (f"S{i} / {iom}" if multi else iom, iom, ip, password)
        for i, (password, ip) in enumerate(shelves, 1)
        for iom in ioms
    ]


# ── 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 for IOM to come back online...")
    time.sleep(30)  # allow time for the IOM to begin shutting down before polling
    if _wait_for_iom_online(password, ip):
        ok(f"{iom} is back online.")
    else:
        warn(f"{iom} did not respond within 5 minutes — proceeding anyway.")
    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 for IOM to come back online...")
    time.sleep(30)  # allow time for the IOM to begin shutting down before polling
    if _wait_for_iom_online(password, ip):
        ok(f"{iom} is back online.")
    else:
        warn(f"{iom} did not respond within 5 minutes — proceeding anyway.")
    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  {_c(C.YEL, '(Do not use on production systems)')}")
    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()
    targets = _make_targets(shelves, iom_choice)

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

    fw_path = _prompt_fw_file("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 iom_choice == "3":
        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 per target (shelf × IOM) ─────────────────────────────────
    for label, iom, ip, password in targets:
        rule(f"{label} ({ip})")
        if update_iom:
            _update_iom_fw(password, ip, iom, fw_path)
        if update_fabric:
            _update_fabric_fw(password, ip, iom, fw_path)

    rule("Post-Update Firmware Validation")
    _show_fw_versions(targets)

    ok("Firmware update complete.")
    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()
    prompt("Press Enter to return to main menu")