"""Auxiliary NCBI tables downloaded separately from the XML corpus.

- journal:           J_Entrez.gz    (~3 MB)  -- complete journal catalog
- deleted_pmid_list: deleted.pmids.gz (~5 MB) -- globally deleted PMIDs
- computed_author:   computed_authors.json (~4 GB) -- disambiguated author clusters

These are small enough (except computed_author) to parse in Python.
computed_author uses DuckDB's streaming read_json for the 4 GB file.
"""
import gzip
import os
import ssl
import sys
from pathlib import Path
from urllib.request import urlopen

import certifi
import duckdb
from coren import MachCap

_SSL = ssl.create_default_context(cafile=certifi.where())

_URLS = {
    "journal": "https://ftp.ncbi.nlm.nih.gov/pubmed/J_Entrez.gz",
    "deleted_pmid_list": "https://ftp.ncbi.nlm.nih.gov/pubmed/deleted.pmids.gz",
    "computed_author": "https://ftp.ncbi.nlm.nih.gov/pub/lu/ComputedAuthors/computed_authors.json",
}


def _download(url: str, dest: Path):
    """Download url to dest atomically with progress."""
    tmp = dest.with_suffix(dest.suffix + '.tmp')
    dest.parent.mkdir(parents=True, exist_ok=True)
    resp = urlopen(url, timeout=120, context=_SSL)
    total = int(resp.headers.get('Content-Length', 0))
    done = 0
    with open(tmp, 'wb') as f:
        while True:
            chunk = resp.read(1 << 18)
            if not chunk:
                break
            f.write(chunk)
            done += len(chunk)
            if total and sys.stderr.isatty():
                print(f"\r  {done*100//total:3d}%  {done>>20}/{total>>20} MB",
                      end='', file=sys.stderr)
    if total and sys.stderr.isatty():
        print(file=sys.stderr)
    os.replace(str(tmp), str(dest))


def _to_parquet(conn, sql: str, output: Path):
    """COPY query result to Parquet atomically."""
    tmp = output.with_suffix('.parquet.tmp')
    conn.execute(f"COPY ({sql}) TO '{tmp}' (FORMAT PARQUET, COMPRESSION ZSTD)")
    os.replace(str(tmp), str(output))


def build_aux_table(data_dir: Path, table_name: str, output: Path):
    """Build one auxiliary Parquet table."""
    _BUILDERS[table_name](data_dir, output)


# -- Journal catalog (J_Entrez.gz) -------------------------------------------

def _build_journal(data_dir: Path, output: Path):
    raw = data_dir / '.aux_cache' / 'J_Entrez.gz'
    print("[pmq] Downloading journal catalog...", file=sys.stderr)
    _download(_URLS["journal"], raw)

    rows = []
    rec: dict[str, str] = {}
    with gzip.open(raw, 'rt', encoding='utf-8', errors='replace') as f:
        for line in f:
            line = line.strip()
            if line.startswith('---'):
                if rec:
                    rows.append((
                        rec.get('NlmId'), rec.get('MedAbbr'),
                        rec.get('JournalTitle'), rec.get('IsoAbbr'),
                        rec.get('ISSN (Print)'), rec.get('ISSN (Online)'),
                    ))
                    rec = {}
                continue
            if ': ' in line:
                k, v = line.split(': ', 1)
                v = v.strip()
                if v:
                    rec[k] = v
        if rec:
            rows.append((
                rec.get('NlmId'), rec.get('MedAbbr'),
                rec.get('JournalTitle'), rec.get('IsoAbbr'),
                rec.get('ISSN (Print)'), rec.get('ISSN (Online)'),
            ))

    conn = duckdb.connect()
    conn.execute("""CREATE TABLE _j(
        nlm_id VARCHAR, medline_ta VARCHAR, title VARCHAR,
        iso_abbr VARCHAR, issn_print VARCHAR, issn_electronic VARCHAR)""")
    conn.executemany("INSERT INTO _j VALUES (?,?,?,?,?,?)", rows)
    _to_parquet(conn, "SELECT * FROM _j ORDER BY nlm_id", output)
    conn.close()
    print(f"[pmq] journal: {len(rows):,} journals", file=sys.stderr)


# -- Deleted PMIDs -----------------------------------------------------------

def _build_deleted_pmids(data_dir: Path, output: Path):
    raw = data_dir / '.aux_cache' / 'deleted.pmids.gz'
    print("[pmq] Downloading deleted PMIDs...", file=sys.stderr)
    _download(_URLS["deleted_pmid_list"], raw)

    pmids = []
    with gzip.open(raw, 'rt') as f:
        for line in f:
            line = line.strip()
            if line.isdigit():
                pmids.append((int(line),))

    conn = duckdb.connect()
    conn.execute("CREATE TABLE _d(pmid INTEGER)")
    conn.executemany("INSERT INTO _d VALUES (?)", pmids)
    _to_parquet(conn, "SELECT * FROM _d ORDER BY pmid", output)
    conn.close()
    print(f"[pmq] deleted_pmid_list: {len(pmids):,} PMIDs", file=sys.stderr)


# -- Computed authors (disambiguated clusters) --------------------------------

def _build_computed_author(data_dir: Path, output: Path):
    """Parse computed_authors.json (~4 GB) via DuckDB streaming read_json.

    Output schema: (name_key, display_name, orcid, pmid)
    One row per author-cluster/PMID pair. Repeated strings are
    dictionary-encoded efficiently in Parquet.
    """
    raw = data_dir / '.aux_cache' / 'computed_authors.json'
    print("[pmq] Downloading computed authors (~4 GB)...", file=sys.stderr)
    _download(_URLS["computed_author"], raw)

    cap = MachCap.read(str(data_dir), detect_net=False, run_bench=False)
    mem_mb = max(256, int(cap.mem_available * 0.5 / (1 << 20)))
    tmp_dir = data_dir / '.duckdb_tmp'
    tmp_dir.mkdir(parents=True, exist_ok=True)

    conn = duckdb.connect()
    conn.execute(f"SET memory_limit='{mem_mb}MB'")
    conn.execute(f"SET temp_directory='{tmp_dir}'")
    conn.execute(f"SET threads={max(1, cap.cores_physical - 1)}")
    conn.execute("SET enable_progress_bar=true")
    conn.execute("SET enable_progress_bar_print=true")

    print("[pmq] Parsing computed authors into Parquet...", file=sys.stderr)
    try:
        _to_parquet(conn, f"""
            SELECT
                name AS name_key,
                names[1] AS display_name,
                CASE WHEN len(orcid) > 0 THEN orcid[1] ELSE NULL END AS orcid,
                UNNEST(pmids) AS pmid
            FROM read_json(
                '{raw}',
                format='newline_delimited',
                columns={{
                    name: 'VARCHAR',
                    names: 'VARCHAR[]',
                    orcid: 'VARCHAR[]',
                    pmids: 'INTEGER[]'
                }}
            )
            ORDER BY pmid
        """, output)
    finally:
        conn.close()

    # Clean up the 4 GB download after successful Parquet creation
    raw.unlink(missing_ok=True)
    sz = output.stat().st_size / (1 << 20)
    print(f"[pmq] computed_author: {sz:.0f} MB", file=sys.stderr)


_BUILDERS = {
    "journal": _build_journal,
    "deleted_pmid_list": _build_deleted_pmids,
    "computed_author": _build_computed_author,
}
