"""asset_photo_extract.py — recover in-cell photos from the Asset File
workbook (DVI-1185 P2).

The workbook's Photo columns hold Excel "Place in cell" rich-value images.
openpyxl surfaces them as #VALUE! errors, but the images themselves live in
xl/media with a mapping chain the OOXML spec spreads over four parts:

    cell @vm (1-based) -> xl/metadata.xml valueMetadata bk[vm-1] rc@v
        -> futureMetadata "XLRICHVALUE" bk[v] xlrd:rvb@i
        -> xl/richData/rdrichvalue.xml rv[i]; the <v> at the structure's
           "_rvRel:LocalImageIdentifier" key position indexes
        -> xl/richData/richValueRel.xml rel[j]@r:id
        -> xl/richData/_rels/richValueRel.xml.rels -> ../media/imageN.ext

This module walks that chain with stdlib only (zipfile + ElementTree) so it
can run server-side against the multi-GB original without openpyxl loading
it, writes each asset's photos to <photo_dir>/<asset_tag>/photoN.ext, and
records them in assets.db via assets_store.set_photos().

CLI:  python3 asset_photo_extract.py <workbook.xlsx> [--dry-run]
      (ASSETS_DB_FILE / ASSETS_PHOTO_DIR env vars override locations)
"""

import argparse
import json
import re
import sys
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path

import assets_store

_NS = {
    "m": "http://schemas.openxmlformats.org/spreadsheetml/2006/main",
    "rd": "http://schemas.microsoft.com/office/spreadsheetml/2017/richdata",
    "rvr": "http://schemas.microsoft.com/office/spreadsheetml/2022/richvaluerel",
    "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
    "pr": "http://schemas.openxmlformats.org/package/2006/relationships",
    "xlrd": "http://schemas.microsoft.com/office/spreadsheetml/2017/richdata",
}

_SKIP_SHEETS = ("Master Data", "Totals")
_CELL_REF_RE = re.compile(r"^([A-Z]+)(\d+)$")


def _col_index(ref):
    """'K3' -> 0-based column index (10)."""
    letters = _CELL_REF_RE.match(ref).group(1)
    n = 0
    for ch in letters:
        n = n * 26 + (ord(ch) - 64)
    return n - 1


def _norm(text):
    return re.sub(r"[^a-z0-9]", "", (text or "").lower())


def _sheet_targets(z):
    """workbook.xml + rels -> [(sheet_name, zip path)] in workbook order."""
    wb = ET.fromstring(z.read("xl/workbook.xml"))
    rels = ET.fromstring(z.read("xl/_rels/workbook.xml.rels"))
    by_id = {rel.get("Id"): rel.get("Target")
             for rel in rels.findall("pr:Relationship", _NS)}
    out = []
    for sheet in wb.find("m:sheets", _NS):
        target = by_id.get(sheet.get(
            "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id"))
        if target:
            out.append((sheet.get("name"), "xl/" + target.lstrip("/")))
    return out


def _shared_strings(z):
    if "xl/sharedStrings.xml" not in z.namelist():
        return []
    root = ET.fromstring(z.read("xl/sharedStrings.xml"))
    return ["".join(t.text or "" for t in si.iter(
        "{http://schemas.openxmlformats.org/spreadsheetml/2006/main}t"))
        for si in root.findall("m:si", _NS)]


def _vm_to_media(z):
    """Precompute cell @vm (1-based) -> media zip path for the whole book."""
    meta = ET.fromstring(z.read("xl/metadata.xml"))
    fm_rvb = []
    for fm in meta.findall("m:futureMetadata", _NS):
        if fm.get("name") != "XLRICHVALUE":
            continue
        for bk in fm.findall("m:bk", _NS):
            rvb = bk.find(".//xlrd:rvb", _NS)
            fm_rvb.append(int(rvb.get("i")) if rvb is not None else None)
    vm_to_fm = [int(bk.find("m:rc", _NS).get("v"))
                for bk in meta.find("m:valueMetadata", _NS).findall("m:bk", _NS)]

    # Key position of the image-relationship index within each structure.
    structs = ET.fromstring(z.read("xl/richData/rdrichvaluestructure.xml"))
    img_key_pos = []
    for s in structs.findall("rd:s", _NS):
        pos = None
        for i, k in enumerate(s.findall("rd:k", _NS)):
            if k.get("n") == "_rvRel:LocalImageIdentifier":
                pos = i
                break
        img_key_pos.append(pos)

    rv_rel_idx = []
    for rv in ET.fromstring(z.read("xl/richData/rdrichvalue.xml")).findall("rd:rv", _NS):
        pos = img_key_pos[int(rv.get("s"))]
        values = rv.findall("rd:v", _NS)
        rv_rel_idx.append(int(values[pos].text)
                          if pos is not None and pos < len(values) else None)

    rel_ids = [rel.get("{http://schemas.openxmlformats.org/officeDocument/2006/"
                       "relationships}id")
               for rel in ET.fromstring(
                   z.read("xl/richData/richValueRel.xml")).findall("rvr:rel", _NS)]
    rel_targets = {rel.get("Id"): "xl/" + rel.get("Target").replace("../", "")
                   for rel in ET.fromstring(
                       z.read("xl/richData/_rels/richValueRel.xml.rels"))
                   .findall("pr:Relationship", _NS)}

    mapping = {}
    for vm_i, fm_i in enumerate(vm_to_fm):
        rvb = fm_rvb[fm_i] if fm_i < len(fm_rvb) else None
        if rvb is None or rvb >= len(rv_rel_idx):
            continue
        rel_i = rv_rel_idx[rvb]
        if rel_i is None or rel_i >= len(rel_ids):
            continue
        media = rel_targets.get(rel_ids[rel_i])
        if media:
            mapping[vm_i + 1] = media
    return mapping


def _sheet_photo_cells(z, path, strings):
    """One sheet -> (photo_cells, header_note). photo_cells is a list of
    (asset_tag, [vm, ...]) per data row that has photo rich values."""
    tag_col = None
    photo_cols = set()
    header_row = None
    rows = []
    for _, row in ET.iterparse(z.open(path), events=("end",)):
        if row.tag != "{http://schemas.openxmlformats.org/spreadsheetml/2006/main}row":
            continue
        row_num = int(row.get("r", "0"))
        cells = {}
        for c in row.findall("m:c", _NS):
            ref = c.get("r")
            if not ref:
                continue
            col = _col_index(ref)
            v = c.find("m:v", _NS)
            text = v.text if v is not None else None
            if c.get("t") == "s" and text is not None:
                text = strings[int(text)]
            cells[col] = {"text": text, "vm": c.get("vm")}
        if header_row is None:
            if row_num <= 5:
                normed = {col: _norm(c["text"]) for col, c in cells.items()}
                if any(h.startswith("assettag") for h in normed.values()):
                    header_row = row_num
                    for col, h in normed.items():
                        if h.startswith("assettag"):
                            tag_col = col
                        elif re.match(r"^photo\d*$", h):
                            photo_cols.add(col)
            row.clear()
            continue
        vms = [int(cells[col]["vm"]) for col in sorted(photo_cols)
               if col in cells and cells[col]["vm"]]
        if vms:
            tag = (cells.get(tag_col, {}).get("text") or "").strip()
            if tag.endswith(".0"):
                tag = tag[:-2]
            rows.append((tag, row_num, vms))
        row.clear()
    return rows, header_row


def extract_photos(workbook_path, dry_run=False):
    """Extract all in-cell photos and attach them to assets.db records.

    Returns a report dict (also stored in assets meta as
    ``last_photo_extract_report``). Photos whose row has no/unknown asset
    tag are flagged, not written."""
    assets_store.init_assets_db()
    photo_dir = assets_store.ASSETS_PHOTO_DIR
    report = {
        "source": str(workbook_path),
        "dry_run": bool(dry_run),
        "assets_with_photos": 0,
        "photos_written": 0,
        "orphan_photos": 0,
        "orphans": [],
        "sheets": {},
    }
    with zipfile.ZipFile(workbook_path) as z:
        vm_media = _vm_to_media(z)
        strings = _shared_strings(z)
        known = set()
        with assets_store._connect() as conn:
            known.update(t for (t,) in conn.execute(
                "SELECT asset_tag FROM assets"))
        for sheet_name, target in _sheet_targets(z):
            if sheet_name in _SKIP_SHEETS:
                continue
            rows, header_row = _sheet_photo_cells(z, target, strings)
            report["sheets"][sheet_name] = sum(len(v) for _, _, v in rows)
            for tag, row_num, vms in rows:
                media = [vm_media[vm] for vm in vms if vm in vm_media]
                if not media:
                    continue
                if not tag or tag not in known:
                    report["orphan_photos"] += len(media)
                    if len(report["orphans"]) < 200:
                        report["orphans"].append(
                            {"sheet": sheet_name, "row": row_num,
                             "tag": tag, "photos": len(media)})
                    continue
                filenames = []
                for i, media_path in enumerate(media, start=1):
                    ext = Path(media_path).suffix or ".img"
                    name = "photo%d%s" % (i, ext)
                    filenames.append(name)
                    if not dry_run:
                        dest = photo_dir / tag
                        dest.mkdir(parents=True, exist_ok=True)
                        (dest / name).write_bytes(z.read(media_path))
                if not dry_run:
                    assets_store.set_photos(tag, filenames)
                report["assets_with_photos"] += 1
                report["photos_written"] += len(filenames)
    if not dry_run:
        assets_store.set_meta("last_photo_extract_report", json.dumps(report))
    return report


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("workbook")
    ap.add_argument("--dry-run", action="store_true",
                    help="map and count only; write nothing")
    args = ap.parse_args(argv)
    report = extract_photos(args.workbook, dry_run=args.dry_run)
    json.dump(report, sys.stdout, indent=1)
    print()
    return 0


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