import hashlib
import json
import ssl
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from pypdf import PdfReader


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "补充公开资料"


def fetch(source):
    result = dict(source)
    destination = OUT / f"{source['id']}_{source['name']}.{source['kind']}"
    try:
        request = urllib.request.Request(source["url"], headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(request, timeout=45, context=ssl.create_default_context()) as response:
            data = response.read()
            result["resolved_url"] = response.url
        if source["kind"] == "pdf" and not data.startswith(b"%PDF"):
            raise ValueError("Expected PDF signature")
        if len(data) < 1000:
            raise ValueError("Response too short")
        destination.write_bytes(data)
        result.update({
            "file": destination.relative_to(ROOT).as_posix(),
            "bytes": len(data),
            "sha256": hashlib.sha256(data).hexdigest(),
            "retrieved": "2026-09-20",
            "download_status": "downloaded",
        })
        if source["kind"] == "pdf":
            reader = PdfReader(destination)
            pages = [{"page": i, "text": (p.extract_text() or "").encode("utf-8", "replace").decode("utf-8")}
                     for i, p in enumerate(reader.pages, 1)]
            result["pages"] = len(pages)
            (ROOT / "research" / "extracted" / f"{source['id']}.json").write_text(
                json.dumps(pages, ensure_ascii=False, indent=2), encoding="utf-8"
            )
            (ROOT / "research" / "extracted" / 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(source["id"], result["download_status"], len(data), flush=True)
    except Exception as error:
        result.update({"download_status": "link_verified_download_pending", "download_error": str(error)})
        print(source["id"], "DOWNLOAD FAILED", str(error), flush=True)
    return result


def main():
    OUT.mkdir(exist_ok=True)
    source_path = ROOT / "research" / "public_sources.json"
    sources = json.loads(source_path.read_text(encoding="utf-8"))
    with ThreadPoolExecutor(max_workers=4) as pool:
        result = list(pool.map(fetch, sources))
    (ROOT / "research" / "public_downloads.json").write_text(
        json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8"
    )


if __name__ == "__main__":
    main()