#!/usr/bin/env python3
"""Extract archive ZIPs and profile every MACStats workbook and worksheet."""

from __future__ import annotations

import argparse
import csv
import json
import re
import zipfile
from pathlib import Path
from typing import Any, Iterable

import openpyxl


EXHIBIT_RE = re.compile(r"(?:EXHIBIT|TABLE)\s*[-_. ]*\s*(\d+[A-Za-z]?)\s*[.:-]?\s*(.*)", re.I)


def read_csv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def normalize_cell(value: Any) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return " ".join(value.replace("\xa0", " ").split())
    return str(value)


def safe_extract(zip_path: Path, destination: Path) -> list[Path]:
    destination.mkdir(parents=True, exist_ok=True)
    extracted: list[Path] = []
    with zipfile.ZipFile(zip_path) as archive:
        for info in archive.infolist():
            member = Path(info.filename.replace("\\", "/"))
            if info.is_dir() or "__MACOSX" in member.parts or member.name.startswith("._"):
                continue
            if member.suffix.lower() not in {".xlsx", ".xlsm", ".xls"}:
                continue
            # Flatten the packaging directory but retain the workbook's name.
            output = destination / member.name
            if output.resolve().parent != destination.resolve():
                raise RuntimeError(f"Unsafe archive member: {info.filename}")
            with archive.open(info) as source, output.open("wb") as target:
                target.write(source.read())
            extracted.append(output)
    return sorted(set(extracted))


def title_from_filename(file_name: str) -> tuple[str, str]:
    stem = Path(file_name).stem.replace("_", " ")
    match = EXHIBIT_RE.search(stem)
    if not match:
        return "", ""
    number = match.group(1)
    title = re.sub(r"[- ]+(?:FY|CY|SFY)s?[- ,].*$", "", match.group(2), flags=re.I)
    title = re.sub(r"[- ]+20\d{2}(?:[-–]\d{2,4})?.*$", "", title)
    title = re.sub(r"[- ]+(?:millions|thousands)(?:\s*\(\d+\))?$", "", title, flags=re.I)
    title = re.sub(r"\s*\(\d+\)$", "", title)
    return number, " ".join(title.replace("-", " ").split())


def find_sheet_title(sample_rows: list[list[str]], sheet_name: str) -> tuple[str, str]:
    candidates: list[str] = []
    for row in sample_rows[:15]:
        text = " ".join(cell for cell in row if cell)
        if text:
            candidates.append(text)
    candidates.append(sheet_name)
    fallback_number = ""
    for candidate in candidates:
        match = EXHIBIT_RE.search(candidate)
        if match:
            fallback_number = fallback_number or match.group(1)
            title = " ".join(match.group(2).split())
            if title:
                return match.group(1), title
    return fallback_number, next((c for c in candidates if len(c) >= 12), sheet_name)


def profile_sheet(ws: openpyxl.worksheet.worksheet.Worksheet, sample_limit: int = 25) -> dict[str, Any]:
    sample_rows: list[list[str]] = []
    nonempty_cells = 0
    last_nonempty_row = 0
    last_nonempty_col = 0
    for row_index, row in enumerate(ws.iter_rows(), start=1):
        rendered = [normalize_cell(cell.value) for cell in row]
        if any(rendered):
            last_nonempty_row = row_index
            last_nonempty_col = max(last_nonempty_col, max(i for i, value in enumerate(rendered, start=1) if value))
            nonempty_cells += sum(bool(value) for value in rendered)
        if row_index <= sample_limit:
            sample_rows.append(rendered[: min(len(rendered), 30)])
    number, title = find_sheet_title(sample_rows, ws.title)
    return {
        "sheet_name": ws.title,
        "sheet_exhibit_number": number,
        "sheet_title": title,
        "max_row": ws.max_row,
        "max_column": ws.max_column,
        "last_nonempty_row": last_nonempty_row,
        "last_nonempty_column": last_nonempty_col,
        "nonempty_cells": nonempty_cells,
        "sample_rows": sample_rows,
    }


def profile_workbook(path: Path) -> tuple[list[dict[str, Any]], str]:
    try:
        workbook = openpyxl.load_workbook(path, read_only=True, data_only=True)
        profiles = [profile_sheet(workbook[name]) for name in workbook.sheetnames]
        workbook.close()
        return profiles, ""
    except Exception as exc:  # retained in QA output for explicit review
        return [], f"{type(exc).__name__}: {exc}"


def write_csv(path: Path, rows: Iterable[dict[str, Any]], fieldnames: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--project-root", type=Path, default=Path(__file__).resolve().parents[1])
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    project = args.project_root
    sources = read_csv(project / "data" / "raw" / "source_manifest.csv")
    extracted_root = project / "data" / "raw" / "extracted"
    workbook_rows: list[dict[str, Any]] = []

    for source in sources:
        if source["excel_available"].lower() not in {"true", "1"}:
            continue
        source_path = Path(source["local_path"])
        if not source_path.is_absolute():
            source_path = project / source_path
        if source["source_kind"] == "archive_zip":
            workbook_paths = safe_extract(source_path, extracted_root / source["report_edition"])
        else:
            workbook_paths = [source_path]
        for workbook_path in workbook_paths:
            file_number, file_title = title_from_filename(workbook_path.name)
            profiles, error = profile_workbook(workbook_path)
            if not profiles:
                workbook_rows.append(
                    {
                        "report_edition": source["report_edition"],
                        "source_kind": source["source_kind"],
                        "workbook_path": str(workbook_path.relative_to(project)),
                        "workbook_file_name": workbook_path.name,
                        "file_exhibit_number": source["exhibit_number"] or file_number,
                        "file_title": source["exhibit_title"] or file_title,
                        "sheet_name": "",
                        "profile_error": error,
                    }
                )
                continue
            for profile in profiles:
                workbook_rows.append(
                    {
                        "report_edition": source["report_edition"],
                        "source_kind": source["source_kind"],
                        "workbook_path": str(workbook_path.relative_to(project)),
                        "workbook_file_name": workbook_path.name,
                        "file_exhibit_number": source["exhibit_number"] or file_number,
                        "file_title": source["exhibit_title"] or file_title,
                        **profile,
                        "profile_error": error,
                    }
                )
            print(f"Profiled {source['report_edition']} / {workbook_path.name} ({len(profiles)} sheets)", flush=True)

    csv_fields = [
        "report_edition", "source_kind", "workbook_path", "workbook_file_name",
        "file_exhibit_number", "file_title", "sheet_name", "sheet_exhibit_number",
        "sheet_title", "max_row", "max_column", "last_nonempty_row",
        "last_nonempty_column", "nonempty_cells", "profile_error",
    ]
    write_csv(project / "data" / "interim" / "workbook_profile.csv", workbook_rows, csv_fields)
    jsonl = project / "data" / "interim" / "workbook_profile.jsonl"
    with jsonl.open("w", encoding="utf-8") as handle:
        for row in workbook_rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    print(f"Wrote {len(workbook_rows)} worksheet profiles")


if __name__ == "__main__":
    main()
