#!/usr/bin/env python
"""
Generate DPM network tables from a Summary sheet.

Usage:
  python build_network_tables.py INPUT.xlsx OUTPUT.xlsx
  python build_network_tables.py INPUT.xlsx OUTPUT.xlsx --net2 DPM1,DPM2,...

Examples:
  python build_network_tables.py MCM08_Network.xlsx MCM08_Network_NET1_NET2.xlsx
  python build_network_tables.py Amazon_CDW5_Bypass.xlsx Bypass_NET1_NET2.xlsx --net2 BYAD_6_DPM1,BYCD_14_DPM1,BYDB_6_DPM1

Requirements: pip install pandas openpyxl
"""

import argparse
import re
from itertools import groupby

import pandas as pd
from openpyxl import Workbook
from openpyxl.styles import PatternFill, Alignment, Border, Side, Font


# ---------- CONSTANTS YOU PROBABLY KEEP THE SAME ----------

SUMMARY_SHEET = "Summary"  # name of sheet with DPM / P_TAG1 / HP

RING_BASE = "11.200.1."       # ring IPs: 11.200.1.2, 11.200.1.3, ...
STAR_NET1_BASE = "11.200.1."  # NET1 STAR base
STAR_NET2_BASE = "11.200.2."  # NET2 STAR base
STAR_START = 20               # first STAR IP is *.20

# VFD HP -> part number
HP_TO_PN = {
    "2HP": "35S-6D2-P101",
    "3HP": "35S-6D3-P101",
    "5HP": "35S-6D4-P111",
    "7.5HP": "35S-6D5-P111",
    "10HP": "35S-6D6-P111",
    "30HP": "25B-D043N114"
}


# ---------- HELPERS ----------

def pick_column(colnames, candidates):
    """Pick first matching column name by case-insensitive list of options."""
    cand_upper = [c.upper() for c in candidates]
    for c in colnames:
        if str(c).strip().upper() in cand_upper:
            return c
    return None


def dpm_sort_key(dpm):
    """Natural sort: prefix (letters/underscore) + list of all numbers."""
    s = str(dpm).strip()
    m = re.match(r"^([A-Za-z_]+)", s)
    prefix = m.group(1) if m else ""
    nums = [int(x) for x in re.findall(r"\d+", s)]
    return prefix, nums


def device_sort_key(tag):
    """
    Sort devices inside a DPM:
    1) FIO Masters/hubs (FIOM/FIOH/FIO/MASTER) first
    2) Then by numbers inside name
    3) VFD before others among non-masters
    """
    s = str(tag).strip()
    u = s.upper()

    is_master = ("FIOM" in u) or ("FIOH" in u) or ("FIO" in u) or ("MASTER" in u)
    master_prio = 0 if is_master else 1

    nums = [int(x) for x in re.findall(r"\d+", s)]

    typ_match = re.search(r"_([A-Za-z]+)\d*$", s)
    typ = typ_match.group(1).upper() if typ_match else ""

    if typ == "VFD":
        t_order = 0
    elif typ in ("FIOM", "FIO", "FIOH"):
        t_order = 1
    else:
        t_order = 2

    return (master_prio, nums, t_order, typ, s)


def part_number_for(device_name, rating_map):
    """Return part number string for a device."""
    if not isinstance(device_name, str):
        return ""

    u = device_name.upper()
    if "SPARE" in u:
        return ""

    # FIO masters / hubs
    if ("FIOM" in u) or ("FIO" in u) or ("FIOH" in u) or ("MASTER" in u):
        return "Murr 54631"

    hp = rating_map.get(device_name, "")
    return HP_TO_PN.get(hp, "")


def build_assignments(df, col_dpm, col_tag, col_rating):
    """Build sorted DPM list, device assignments and rating map."""
    df = df[~df[col_dpm].isna()].copy()

    unique_dpms = sorted(
        {str(x).strip() for x in df[col_dpm].tolist()},
        key=dpm_sort_key,
    )

    assign = {dpm: [] for dpm in unique_dpms}
    rating_map = {}

    for _, row in df.iterrows():
        dpm = str(row[col_dpm]).strip()
        tag = (
            str(row[col_tag]).strip()
            if col_tag and not pd.isna(row[col_tag])
            else None
        )
        rating = (
            str(row[col_rating]).strip()
            if col_rating and not pd.isna(row[col_rating])
            else ""
        )
        if tag:
            assign.setdefault(dpm, []).append(tag)
            rating_map[tag] = rating

    for dpm in assign:
        assign[dpm] = sorted(assign[dpm], key=device_sort_key)

    return unique_dpms, assign, rating_map


def generate_rows(unique_dpms, assign, rating_map, net2_dpms):
    """
    Generate rows for NET1 and NET2.
    Each row: (dpm, ring_ip, assigned_dev, part_num, star_ip, dpm_port)
    """
    # ring IP per DPM
    ring_ip_map = {}
    ring_counter = 2
    for dpm in unique_dpms:
        ring_ip_map[dpm] = f"{RING_BASE}{ring_counter}"
        ring_counter += 1

    rows_net1 = []
    rows_net2 = []

    star1_counter = STAR_START
    star2_counter = STAR_START

    for dpm in unique_dpms:
        ring_ip = ring_ip_map[dpm]
        tags = assign.get(dpm, [])
        
        # Generate ports starting from 5, continuing beyond 28 if needed
        num_devices = len(tags)
        ports = list(range(5, 5 + num_devices))

        for idx, port in enumerate(ports):
            dev = tags[idx]
            pn = part_number_for(dev, rating_map)

            if dpm in net2_dpms:
                star_ip = f"{STAR_NET2_BASE}{star2_counter}"
                star2_counter += 1
                rows_net2.append(
                    (dpm, ring_ip, dev, pn, star_ip, f"{dpm}_P{port}")
                )
            else:
                star_ip = f"{STAR_NET1_BASE}{star1_counter}"
                star1_counter += 1
                rows_net1.append(
                    (dpm, ring_ip, dev, pn, star_ip, f"{dpm}_P{port}")
                )

    return rows_net1, rows_net2


def write_sheet(ws, rows, title):
    """Write one NET sheet with merged yellow bands, D..I layout."""
    ws.title = title

    headers = ["DPM", "IP", "Assigned Device", "Part Number", "IP", "DPM PORT"]
    for col, h in enumerate(headers, start=4):
        c = ws.cell(row=3, column=col, value=h)
        c.font = Font(bold=True)
        c.alignment = Alignment(horizontal="center", vertical="center")
        c.border = Border(
            left=Side(style="thin"),
            right=Side(style="thin"),
            top=Side(style="thin"),
            bottom=Side(style="thin"),
        )

    yellow = PatternFill("solid", fgColor="FFC000")
    red = PatternFill("solid", fgColor="FF0000")
    thin = Side(style="thin", color="000000")
    border = Border(left=thin, right=thin, top=thin, bottom=thin)
    center = Alignment(horizontal="center", vertical="center")
    left = Alignment(horizontal="left", vertical="center")

    widths = {4: 18, 5: 16, 6: 30, 7: 18, 8: 16, 9: 22}
    for col, w in widths.items():
        ws.column_dimensions[chr(64 + col)].width = w

    row_idx = 4

    if not rows:
        ws.freeze_panes = "D4"
        return

    for dpm, group in groupby(rows, key=lambda x: x[0]):
        g = list(group)
        num_rows = len(g)

        start = row_idx
        end = row_idx + num_rows - 1

        ws.merge_cells(start_row=start, start_column=4, end_row=end, end_column=4)
        ws.merge_cells(start_row=start, start_column=5, end_row=end, end_column=5)

        cD = ws.cell(row=start, column=4, value=g[0][0])
        cE = ws.cell(row=start, column=5, value=g[0][1])
        for cell in (cD, cE):
            cell.fill = yellow
            cell.alignment = center
            cell.border = border
            cell.font = Font(bold=False)

        for i in range(num_rows):
            row_vals = {
                6: g[i][2],  # Assigned Device
                7: g[i][3],  # Part Number
                8: g[i][4],  # IP
                9: g[i][5],  # DPM PORT
            }
            # Highlight rows beyond 24th device in red
            is_overflow = i >= 24
            fill_color = red if is_overflow else None
            
            for col, val in row_vals.items():
                c = ws.cell(row=start + i, column=col, value=val)
                c.alignment = left if col in (6, 7, 9) else center
                c.border = border
                if fill_color:
                    c.fill = fill_color

        for r in range(start + 1, end + 1):
            for col in (4, 5):
                c = ws.cell(row=r, column=col)
                c.fill = yellow
                c.border = border

        row_idx = end + 1

    ws.freeze_panes = "D4"


def main():
    parser = argparse.ArgumentParser(description="Build DPM network tables.")
    parser.add_argument("input_file", help="Input Excel file (with Summary sheet)")
    parser.add_argument("output_file", help="Output Excel file")
    parser.add_argument(
        "--net2",
        help="Comma-separated list of DPM names that should go to NET2 (.2 subnet)",
        default="",
    )
    args = parser.parse_args()

    net2_dpms = {
        x.strip()
        for x in args.net2.split(",")
        if x.strip()
    }

    # Load Summary
    df = pd.read_excel(args.input_file, sheet_name=SUMMARY_SHEET)
    df.columns = [str(c).strip() for c in df.columns]

    col_dpm = pick_column(df.columns, ["DPM"])
    col_tag = pick_column(df.columns, ["P_TAG1", "P TAG1", "ASSIGNED DEVICE"])
    col_rating = pick_column(df.columns, ["RATING", "HP", "POWER", "RATED"])

    if not col_dpm or not col_tag:
        raise RuntimeError(
            f"Could not find DPM or P_TAG1/Assigned Device columns in sheet '{SUMMARY_SHEET}'"
        )

    unique_dpms, assign, rating_map = build_assignments(
        df, col_dpm, col_tag, col_rating
    )

    rows_net1, rows_net2 = generate_rows(unique_dpms, assign, rating_map, net2_dpms)

    # Build workbook
    wb = Workbook()
    ws1 = wb.active
    write_sheet(ws1, rows_net1, "NET1 STRUCTURE")

    ws2 = wb.create_sheet()
    write_sheet(ws2, rows_net2, "NET2 STRUCTURE")

    wb.save(args.output_file)
    print(f"Saved: {args.output_file}")


if __name__ == "__main__":
    main()