#!/usr/bin/env python3
"""Fill AIME production order template from Odoo MRP export.

Usage:
    python fill_production_order.py <source_xlsx>
    python fill_production_order.py               # uses newest matching source file

The script loads the specified Excel source file, extracts relevant fields,
fills the ODS template, and saves it as a new ODS named after cell B1:C2
of the resulting template (MODEL-FA without the prefixes).
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

import openpyxl
from odf import opendocument
from odf.table import Table, TableCell, TableRow
from odf.text import P

# ---------------------------------------------------------------------------
# Paths / globs
# ---------------------------------------------------------------------------
TEMPLATE_GLOB = "A4004_*.ods"
SOURCE_GLOB = "Production Order (mrp.production)*.xlsx"

# ---------------------------------------------------------------------------
# Source column indices (0-based).  Column D = index 3, E = 4, G = 6.
# ---------------------------------------------------------------------------
COL_CAT = 3   # D
COL_NAME = 4  # E
COL_SN = 6    # G

# ---------------------------------------------------------------------------
# Mapping from source component category to target template rows.
# Template columns (1-based as the user sees them):
#   Category = B (col 1), Name = C:D (cols 2-3), Serial = E (col 4).
# The rows are fixed by the template layout.
# ---------------------------------------------------------------------------
CATEGORY_TARGETS: dict[str, dict] = {
    "All / Components / Barebone": {
        "name_row": 10,
        "serial_row": 10,
        "max_serials": 1,
    },
    "All / Components / CPU": {
        "name_row": 11,
        "serial_row": 11,
        "max_serials": 1,
    },
    "All / Components / RAM": {
        "name_row": 12,
        "serial_row": 12,
        "max_serials": 8,
    },
    "All / Components / GPU": {
        "name_row": 20,
        "serial_row": 20,
        "max_serials": 4,
    },
    "All / Components / SSD": {
        "name_row": 24,
        "serial_row": 24,
        "max_serials": 3,
    },
    "All / Components / Network": {
        "name_row": 30,
        "serial_row": 30,  # NOTE: user's spec says E30 even though table row is 30
        "max_serials": 2,
    },
    "All / Components / RAID": {
        "name_row": 32,
        "serial_row": 32,  # NOTE: user's spec says E30; see fill logic below
        "max_serials": 1,
    },
}

# Order in which categories are processed.
CATEGORY_ORDER = [
    "All / Components / Barebone",
    "All / Components / CPU",
    "All / Components / RAM",
    "All / Components / GPU",
    "All / Components / SSD",
    "All / Components / Network",
    "All / Components / RAID",
]

# Source categories that may appear without an explicit category cell but belong
# to the preceding known category (e.g. continuation rows for RAM / SSD).

# ---------------------------------------------------------------------------
# ODS helpers
# ---------------------------------------------------------------------------
def _get_rows(table: Table) -> list[TableRow]:
    return list(table.getElementsByType(TableRow))


def _get_cells(row: TableRow) -> list[TableCell]:
    """Return all real ``table:table-cell`` elements of ``row`` in document order.

    ``table:covered-table-cell`` placeholders are *not* returned because they
    are only merge markers; they do not represent distinct visual columns.
    """
    cells: list[TableCell] = []
    for child in row.childNodes:
        if hasattr(child, "tagName") and child.tagName == "table:table-cell":
            cells.append(child)
    return cells


def _physical_column_index(row: TableRow, col: int) -> tuple[int, int]:
    """Map a logical/visual column index to a physical child in ``row``.

    Returns ``(physical_index, logical_start)`` where ``logical_start`` is the
    first logical column covered by the returned physical cell.

    A logical column is what the user sees (A=0, B=1, C=2, ...).  In ODF a
    merged cell is represented by one ``table:table-cell`` with span and
    repetition attributes followed by ``table:covered-table-cell``
    placeholders.  A covered placeholder directly following a merge owner in
    the same row is only a marker and contributes no visual width.  A covered
    placeholder standing on its own (e.g. in a continuation row of a vertical
    merge) represents the merged area and contributes its own span width.
    """
    logical = 0
    merge_marker_count = 0
    for physical, child in enumerate(row.childNodes):
        if not hasattr(child, "tagName"):
            continue
        tag = child.tagName
        if tag == "table:covered-table-cell":
            span = int(
                child.attributes.get(
                    ("urn:oasis:names:tc:opendocument:xmlns:table:1.0", "number-columns-spanned"),
                    1,
                )
            )
            repeat = int(
                child.attributes.get(
                    ("urn:oasis:names:tc:opendocument:xmlns:table:1.0", "number-columns-repeated"),
                    1,
                )
            )
            if merge_marker_count > 0:
                merge_marker_count -= 1
                # This placeholder is part of the preceding horizontal merge.
                continue
            width = span * repeat
        elif tag == "table:table-cell":
            span = int(
                child.attributes.get(
                    ("urn:oasis:names:tc:opendocument:xmlns:table:1.0", "number-columns-spanned"),
                    1,
                )
            )
            repeat = int(
                child.attributes.get(
                    ("urn:oasis:names:tc:opendocument:xmlns:table:1.0", "number-columns-repeated"),
                    1,
                )
            )
            # The merge owner itself already accounts for all of its columns.
            # Any covered placeholders directly following it in this row are
            # markers for the covered part of the merge.
            merge_marker_count = max(0, span - 1)
            width = span * repeat
        else:
            continue

        if logical <= col < logical + width:
            return physical, logical
        logical += width
    raise IndexError(f"Logical column {col} not found in row")


def _clone_empty_cell(template: TableCell) -> TableCell:
    """Create an empty cell with the same style/span as ``template``."""
    cell = TableCell()
    style = template.getAttribute("stylename")
    if style:
        cell.setAttribute("stylename", style)
    for attr in ("numbercolumnsspanned", "numberrowsspanned"):
        val = template.getAttribute(attr)
        if val is not None:
            cell.setAttribute(attr, val)
    cell.addElement(P(text=""))
    return cell


def _ensure_rows(table: Table, min_row_count: int) -> None:
    """Append empty rows until the table has at least ``min_row_count`` rows."""
    rows = _get_rows(table)
    if not rows:
        return
    first_row_cells = _get_cells(rows[0])
    while len(rows) < min_row_count:
        new_row = TableRow()
        for _ in first_row_cells:
            cell = TableCell()
            cell.addElement(P(text=""))
            new_row.addElement(cell)
        table.addElement(new_row)
        rows = _get_rows(table)


def _normalize_row_heights(table: Table, height_in: str = "0.18in") -> None:
    """Set all rows in ``table`` to the same fixed height.

    The template's automatic row styles often use ``use-optimal-row-height``,
    which causes rows with multi-line merged cells (e.g. the RAM continuation
    rows) to grow unnecessarily tall.  This function replaces each row's
    style reference with a fresh automatic row style that has a fixed height.
    """
    from odf.style import Style
    from odf.table import TableRowProperties

    doc = table.parentNode
    while doc is not None and not hasattr(doc, "automaticstyles"):
        doc = doc.parentNode
    if doc is None:
        return

    for idx, row in enumerate(_get_rows(table)):
        style_name = f"rowheight{idx}"
        style = Style(name=style_name, family="table-row")
        props = TableRowProperties()
        props.setAttribute("rowheight", height_in)
        props.setAttribute("breakbefore", "auto")
        style.addElement(props)
        doc.automaticstyles.addElement(style)
        row.setAttribute("stylename", style_name)


def _set_cell_value(table: Table, row: int, col: int, value: str) -> None:
    """Set the text value of the cell at logical ``row``/``col`` (0-based).

    ``col`` is a logical/visual column index (A=0, B=1, C=2, ...).  Merged
    cells are mapped to the correct physical child automatically.  When the
    target physical cell is a ``covered-table-cell`` placeholder inside a
    merge it is converted to a regular ``table-cell`` so the value becomes
    visible and editable.
    """
    _ensure_rows(table, row + 1)
    rows = _get_rows(table)
    target_row = rows[row]
    physical_col, logical_start = _physical_column_index(target_row, col)

    row_children = list(target_row.childNodes)
    if physical_col >= len(row_children):
        raise IndexError(
            f"Row {row + 1} only has {len(row_children)} child nodes, cannot write to logical column {col + 1}"
        )

    cell = row_children[physical_col]
    style = cell.getAttribute("stylename")

    # Preserve merge attributes.
    preserved_spans = {
        attr: cell.getAttribute(attr)
        for attr in ("numbercolumnsspanned", "numberrowsspanned")
        if cell.getAttribute(attr) is not None
    }

    # If the target cell is repeated, split it so we only modify the exact
    # logical column requested.  Otherwise odfpy may place the value in the
    # wrong repeated instance.
    repeat = int(
        cell.attributes.get(
            ("urn:oasis:names:tc:opendocument:xmlns:table:1.0", "number-columns-repeated"),
            1,
        )
    )
    span = int(
        cell.attributes.get(
            ("urn:oasis:names:tc:opendocument:xmlns:table:1.0", "number-columns-spanned"),
            1,
        )
    )
    if repeat > 1:
        offset = (col - logical_start) // span
        row_children = list(target_row.childNodes)
        physical_idx_in_row = row_children.index(cell)
        target_row.removeChild(cell)

        # Insert empty cells for repeats before the target instance.
        if offset > 0:
            before = _clone_empty_cell(cell)
            before.setAttribute("numbercolumnsrepeated", str(offset))
            target_row.insertBefore(before, row_children[physical_idx_in_row + 1] if physical_idx_in_row + 1 < len(row_children) else None)

        # The target instance becomes a single non-repeated cell.
        cell = TableCell()
        if style:
            cell.setAttribute("stylename", style)
        for attr, val in preserved_spans.items():
            cell.setAttribute(attr, val)
        target_row.insertBefore(cell, row_children[physical_idx_in_row + 1] if physical_idx_in_row + 1 < len(row_children) else None)

        # Insert empty cells for repeats after the target instance.
        remaining = repeat - offset - 1
        if remaining > 0:
            after = _clone_empty_cell(cell)
            after.setAttribute("numbercolumnsrepeated", str(remaining))
            target_row.insertBefore(after, row_children[physical_idx_in_row + 1] if physical_idx_in_row + 1 < len(row_children) else None)

    # Clear existing paragraphs.
    for child in list(cell.childNodes):
        cell.removeChild(child)

    if style:
        cell.setAttribute("stylename", style)

    for attr, val in preserved_spans.items():
        cell.setAttribute(attr, val)

    cell.addElement(P(text=str(value) if value is not None else ""))


def _find_template(directory: Path) -> Path:
    candidates = sorted(directory.glob(TEMPLATE_GLOB))
    if not candidates:
        raise FileNotFoundError(
            f"No ODS template matching '{TEMPLATE_GLOB}' found in {directory}"
        )
    return candidates[0]


# ---------------------------------------------------------------------------
# Source parsing
# ---------------------------------------------------------------------------
def _clean(value) -> str | None:
    if value is None:
        return None
    text = str(value).strip()
    return text if text else None


def _count_empty_name_rows(ws, start_row: int) -> tuple[int, list[str]]:
    """Count consecutive continuation rows and collect their serial numbers.

    Starting on the row *below* ``start_row`` we walk downwards while both
    the category cell (column D) and the name cell (column E) are empty.
    For every such row that has a serial number in column G we collect it.
    We stop as soon as a non-empty category or name appears.

    Returns (number_of_continuation_rows, list_of_extra_serials_found).
    """
    extra_serials: list[str] = []
    max_row = ws.max_row
    r = start_row + 1
    while r <= max_row:
        cat_val = _clean(ws.cell(row=r, column=COL_CAT + 1).value)
        name_val = _clean(ws.cell(row=r, column=COL_NAME + 1).value)
        if cat_val is not None or name_val is not None:
            break
        sn = _clean(ws.cell(row=r, column=COL_SN + 1).value)
        if sn:
            extra_serials.append(sn)
        r += 1
    return (r - start_row - 1), extra_serials


def _extract_blocks(ws) -> dict[str, list[dict]]:
    """Extract every component block from the source workbook.

    Returns a dict mapping source category -> list of blocks.
    Each block has:
        - row:       first source row of the block
        - name:      component name (column E)
        - serials:   list of serial numbers belonging to this block

    A new block starts when a known category cell appears or when a non-empty
    name cell appears directly below a known category that did not yet have a
    name in its own row.
    """
    blocks: dict[str, list[dict]] = {cat: [] for cat in CATEGORY_TARGETS}

    max_row = ws.max_row
    r = 1
    while r <= max_row:
        category = _clean(ws.cell(row=r, column=COL_CAT + 1).value)
        name = _clean(ws.cell(row=r, column=COL_NAME + 1).value)
        serial = _clean(ws.cell(row=r, column=COL_SN + 1).value)

        if category not in CATEGORY_TARGETS:
            r += 1
            continue

        # Skip truly empty category rows.
        if name is None and serial is None:
            r += 1
            continue

        serials: list[str] = []
        if serial:
            serials.append(serial)

        _, extra_serials = _count_empty_name_rows(ws, r)
        serials.extend(extra_serials)

        blocks[category].append({
            "row": r,
            "name": name,
            "serials": serials,
        })

        r += 1 + len(extra_serials)

    return blocks


def extract_source_data(source_path: Path) -> dict:
    """Read the source workbook and return extracted values."""
    wb = openpyxl.load_workbook(source_path, data_only=True)
    ws = wb.active

    # B2 -> FA, C2 -> MODEL
    b2 = _clean(ws.cell(row=2, column=2).value)
    c2 = _clean(ws.cell(row=2, column=3).value)

    fa: str | None = None
    model: str | None = None
    if isinstance(b2, str):
        m = re.search(r"FA-([A-Za-z0-9]+)", b2)
        if m:
            fa = m.group(1)
    if isinstance(c2, str):
        m = re.search(r"AIME-([A-Za-z0-9]+)", c2)
        if m:
            model = m.group(1)

    if fa is None:
        raise ValueError(f"Could not find 'FA-' serial number in {source_path}")
    if model is None:
        raise ValueError(f"Could not find 'AIME-' model in {source_path}")

    blocks = _extract_blocks(ws)
    wb.close()

    full_model = f"{model}-{fa}"

    components: dict[str, dict] = {cat: {"name": None, "serials": []} for cat in CATEGORY_TARGETS}
    ssd_second: dict | None = None

    counters = {
        "mGPU": 0,
        "mRAM": 0,
        "mSSD": 0,
        "mNIC": 0,
    }

    # Per the user's variable definition, the counters count how many *empty*
    # name rows exist below the component before the next non-empty name.
    # Those empty rows are the rows that contain extra serial numbers.

    for cat in CATEGORY_ORDER:
        cat_blocks = blocks.get(cat, [])
        if not cat_blocks:
            continue

        first = cat_blocks[0]

        # SSD handling: the example output places the first SSD name in the
        # primary slot (row 24) and the second distinct SSD name in the
        # secondary slot (rows 27-29).  If the primary block has no serials
        # but a later block does, borrow serials from the last block so the
        # primary slot is filled (matching the supplied example file).
        if cat == "All / Components / SSD" and len(cat_blocks) > 1:
            merged_serials = list(first["serials"])
            ssd_second_idx: int | None = None
            for idx, blk in enumerate(cat_blocks[1:], start=1):
                if blk["name"] == first["name"]:
                    merged_serials.extend(blk["serials"])
                elif ssd_second_idx is None:
                    ssd_second_idx = idx
            first = {
                **first,
                "serials": merged_serials,
            }
            if ssd_second_idx is not None:
                ssd_second = cat_blocks[ssd_second_idx]
            # If the primary block ended up without serials, borrow from the
            # last SSD block so the primary slot is not left empty.
            if not first["serials"] and cat_blocks:
                first = {
                    **first,
                    "serials": list(cat_blocks[-1]["serials"]),
                }

        components[cat]["name"] = first["name"]
        components[cat]["serials"] = first["serials"]

        # The m* variables count how many *extra* serial rows exist beyond the first.
        counter_key = {
            "All / Components / GPU": "mGPU",
            "All / Components / RAM": "mRAM",
            "All / Components / SSD": "mSSD",
            "All / Components / Network": "mNIC",
        }.get(cat)
        if counter_key:
            counters[counter_key] = max(0, len(first["serials"]) - 1)

    return {
        "fa": fa,
        "model": model,
        "full_model": full_model,
        "components": components,
        "ssd_second": ssd_second,
        **counters,
    }


# ---------------------------------------------------------------------------
# Template filling
# ---------------------------------------------------------------------------
def fill_template(source_path: Path, output_path: Path | None = None) -> Path:
    directory = source_path.parent
    data = extract_source_data(source_path)

    template_path = _find_template(directory)
    doc = opendocument.load(template_path)

    tables = list(doc.spreadsheet.getElementsByType(Table))
    if not tables:
        raise ValueError("No table found in the template ODS file")
    table = tables[0]

    print(f"Source: {source_path.name}")
    print(f"FA: {data['fa']}, MODEL: {data['model']}, Combined: {data['full_model']}")
    print(
        f"Counters -> mGPU={data['mGPU']}, mRAM={data['mRAM']}, "
        f"mSSD={data['mSSD']}, mNIC={data['mNIC']}"
    )

    # Combined model-serial into B1 (merged B1:C2).
    _set_cell_value(table, 0, 1, data["full_model"])

    # Logical column indices (A=0, B=1, C=2, ...).
    name_col = 2   # C (merge owner for the C:D name cell)
    serial_col = 4 # E / S/N column

    for category in CATEGORY_ORDER:
        comp = data["components"][category]
        targets = CATEGORY_TARGETS[category]
        name = comp.get("name")
        serials = comp.get("serials", [])

        if name:
            row_idx = targets["name_row"] - 1
            _set_cell_value(table, row_idx, name_col, name)
            print(f"  {category} name -> C{targets['name_row']}:D{targets['name_row']}: {name}")

        for i, sn in enumerate(serials[: targets["max_serials"]]):
            target_row = targets["serial_row"] - 1 + i
            _set_cell_value(table, target_row, serial_col, sn)
            print(f"  {category} S/N {i + 1} -> E{target_row + 1}: {sn}")


    # Secondary SSD block into rows 27-29 (name C27:D27, serial E27:E29).
    ssd_second = data.get("ssd_second")
    if ssd_second:
        name = ssd_second.get("name")
        serials = ssd_second.get("serials", [])
        print(f"DEBUG ssd_second name={name!r} serials={serials}")
        if name:
            _set_cell_value(table, 26, name_col, name)
            print(f"  Second SSD name -> C27:D27: {name}")
        for i, sn in enumerate(serials[:3]):
            _set_cell_value(table, 26 + i, serial_col, sn)
            print(f"  Second SSD S/N {i + 1} -> E{27 + i}: {sn}")

    # Force a uniform row height so multi-line merged cells do not expand rows.
    _normalize_row_heights(table)

    if output_path is None:
        output_path = directory / f"{data['full_model']}.ods"
    else:
        output_path = Path(output_path)

    doc.save(output_path)
    print(f"Saved: {output_path}")
    return output_path


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Fill production order template from an Odoo export."
    )
    parser.add_argument(
        "source",
        nargs="?",
        help="Path to the source Excel file. If omitted, the alphabetically last "
        "source file matching the production order glob is used.",
    )
    parser.add_argument(
        "-o", "--output", help="Optional explicit output ODS file path"
    )
    args = parser.parse_args(argv)

    directory = Path(__file__).resolve().parent

    if args.source:
        source_path = Path(args.source).expanduser()
        if not source_path.is_absolute():
            source_path = directory / source_path
    else:
        sources = sorted(directory.glob(SOURCE_GLOB))
        if not sources:
            print(f"No source file matching '{SOURCE_GLOB}' found.", file=sys.stderr)
            return 1
        source_path = sources[-1]
        print(f"Using source file: {source_path.name}")

    if not source_path.exists():
        print(f"Source file not found: {source_path}", file=sys.stderr)
        return 1

    fill_template(source_path, args.output)
    return 0


if __name__ == "__main__":
    sys.exit(main())
