import hashlib
import json
import re
import zipfile
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from xml.etree import ElementTree as ET

from pypdf import PdfReader


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "research"
TEXT = OUT / "extracted"
NS = {"s": "http://schemas.openxmlformats.org/spreadsheetml/2006/main"}


def workbook_rows(path):
    result = {"file": path.name, "sheets": []}
    with zipfile.ZipFile(path) as archive:
        strings = []
        if "xl/sharedStrings.xml" in archive.namelist():
            tree = ET.fromstring(archive.read("xl/sharedStrings.xml"))
            strings = ["".join(node.itertext()) for node in tree.findall("s:si", NS)]
        workbook = ET.fromstring(archive.read("xl/workbook.xml"))
        reltree = ET.fromstring(archive.read("xl/_rels/workbook.xml.rels"))
        rels = {node.attrib["Id"]: node.attrib["Target"] for node in reltree}
        for sheet in workbook.findall("s:sheets/s:sheet", NS):
            rid = sheet.attrib["{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id"]
            target = rels[rid]
            member = target.lstrip("/") if target.startswith("/") else "xl/" + target
            rows = []
            with archive.open(member) as stream:
                for _, node in ET.iterparse(stream, events=("end",)):
                    if node.tag != f"{{{NS['s']}}}row":
                        continue
                    cells = {}
                    for cell in node.findall("s:c", NS):
                        value = cell.find("s:v", NS)
                        text = value.text if value is not None else None
                        kind = cell.attrib.get("t")
                        if kind == "s" and text is not None:
                            text = strings[int(text)]
                        elif kind == "inlineStr":
                            text = "".join(cell.find("s:is", NS).itertext())
                        if text is not None and str(text).strip():
                            cells[cell.attrib["r"]] = text
                    if cells:
                        rows.append({"row": int(node.attrib["r"]), "cells": cells})
                    node.clear()
            result["sheets"].append({"name": sheet.attrib["name"], "rows": rows})
    return result


def pdf_record(item):
    index, path = item
    source_id = f"L{index:02}"
    raw = path.read_bytes()
    reader = PdfReader(path)
    pages = []
    for number, page in enumerate(reader.pages, 1):
        text = (page.extract_text() or "").encode("utf-8", errors="replace").decode("utf-8")
        pages.append({"page": number, "text": text})
    metadata = {str(key): str(value) for key, value in (reader.metadata or {}).items()}
    record = {
        "id": source_id,
        "file": path.relative_to(ROOT).as_posix(),
        "name": path.name,
        "sha256": hashlib.sha256(raw).hexdigest(),
        "bytes": len(raw),
        "pages": len(pages),
        "text_pages": sum(len(re.sub(r"\s", "", p["text"])) > 80 for p in pages),
        "metadata": metadata,
        "cover": "\n".join(p["text"] for p in pages[:2])[:3500],
    }
    (TEXT / f"{source_id}.json").write_text(
        json.dumps(pages, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    (TEXT / f"{source_id}.txt").write_text(
        "\n\n".join(f"--- PDF PAGE {p['page']} ---\n{p['text']}" for p in pages),
        encoding="utf-8",
    )
    print(f"{source_id}\t{len(pages)}\t{record['text_pages']}\t{path.name}", flush=True)
    return record


def main():
    TEXT.mkdir(parents=True, exist_ok=True)
    folders = [
        ROOT / "厨房用具-行业洞察报告-何英杰-2026.9.18",
        ROOT / "水杯茶杯水壶保温杯行业报告",
    ]
    paths = sorted(p for folder in folders for p in folder.rglob("*.pdf"))
    with ThreadPoolExecutor(max_workers=4) as pool:
        records = list(pool.map(pdf_record, enumerate(paths, 1)))
    first_seen = {}
    for record in records:
        digest = record["sha256"]
        record["duplicate_of"] = first_seen.get(digest)
        first_seen.setdefault(digest, record["id"])
    (OUT / "local_sources.json").write_text(
        json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    workbooks = [workbook_rows(path) for path in sorted(ROOT.glob("*.xlsx"))]
    (OUT / "workbooks.json").write_text(
        json.dumps(workbooks, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(json.dumps({
        "pdf_files": len(records),
        "unique_pdfs": len(first_seen),
        "pages": sum(r["pages"] for r in records),
        "workbooks": [
            {"file": w["file"], "sheets": [
                {"name": s["name"], "populated_rows": len(s["rows"])}
                for s in w["sheets"]
            ]} for w in workbooks
        ],
    }, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()