"""pmq -- PubMed database builder and query interface.

    from pmq import PubMed

    with PubMed() as db:
        db.df("SELECT * FROM citation WHERE pub_year = '2024' LIMIT 10")

`python -m pmq` builds the dataset, compacts tables, and starts serving
via irohds P2P. Every machine that builds also serves, so new users
get instant downloads instead of multi-hour builds.
"""
__version__ = "0.0.6"
import os
import sys
import time
from pathlib import Path

import duckdb
import irohds
from coren import MachCap

from ._build import default_data_dir, write_build_date, ensure_build
from ._tables import create_table, ftp_fingerprint, has_all_tables, table_has_data, ALL_TABLES

_MIN_DISK_BYTES = 40 << 30  # 40 GB


def _check_resources(data_dir: Path):
    cap = MachCap.read(str(data_dir), detect_net=False, run_bench=False)
    if cap.disk_free < _MIN_DISK_BYTES:
        gb_free = cap.disk_free / (1 << 30)
        gb_need = _MIN_DISK_BYTES / (1 << 30)
        print(
            f"[pmq] WARNING: only {gb_free:.1f} GB free on {data_dir}. "
            f"Need ~{gb_need:.0f} GB for Parquet output.\n"
            f"[pmq] Use PubMed(data_dir='/path/to/bigger/drive') or "
            f"symlink: ln -s /mnt/ssd/pmq {data_dir}",
            file=sys.stderr,
        )


def _compact_and_serve(data_dir: Path, serve: bool = True):
    """Compact tables and register with irohds."""
    os.environ["IROHDS_DATA_DIR"] = str(data_dir)
    fp = ftp_fingerprint()
    n = len(ALL_TABLES)
    t0 = time.time()
    print(f"[pmq] Verifying {n} tables...", file=sys.stderr)
    for i, name in enumerate(ALL_TABLES, 1):
        if not table_has_data(data_dir, name):
            print(f"  [{i:2d}/{n}] {name:25s}       -- MB  empty", file=sys.stderr)
            continue
        t1 = time.time()
        ref = create_table(fp, name, data_dir=str(data_dir))
        elapsed = time.time() - t1
        size_mb = os.path.getsize(ref.path) / (1 << 20) if os.path.exists(ref.path) else 0
        status = "compact" if elapsed > 2.0 else "ok"
        print(f"  [{i:2d}/{n}] {name:25s} {size_mb:7.1f} MB  {status} ({elapsed:.1f}s)",
              file=sys.stderr)
    total = time.time() - t0
    print(f"[pmq] Tables ready in {total:.0f}s", file=sys.stderr)
    if has_all_tables(data_dir):
        write_build_date(data_dir)
    if serve:
        try:
            irohds.status("pmq")
            print("[pmq] Serving on irohds namespace 'pmq'.", file=sys.stderr)
        except Exception as e:
            print(f"[pmq] WARNING: irohds status: {e}", file=sys.stderr)


class PubMed:
    """Query interface over PubMed Parquet tables.

    On first query, create_table() is called for each table via irohds.
    The memo key ties code + parameters to data: if the FTP listing
    changes, tables are automatically re-compacted. If unchanged, the
    cached result is returned (from local cache or P2P).
    """

    def __init__(self, data_dir: str | None = None, *, through: str | None = None):
        self.data_dir = Path(data_dir).resolve() if data_dir else default_data_dir()
        self._through = through
        self._conn: duckdb.DuckDBPyConnection | None = None
        self._views = False

    def __enter__(self):
        os.environ["IROHDS_DATA_DIR"] = str(self.data_dir)
        _check_resources(self.data_dir)
        self._conn = duckdb.connect()
        return self

    def __exit__(self, *exc):
        if self._conn:
            self._conn.close()
            self._conn = None

    def _ensure_views(self):
        if self._views:
            return
        fp = ftp_fingerprint()
        dd = str(self.data_dir)
        year_filter = ""
        if self._through:
            year = str(self._through)[:4]
            year_filter = f" WHERE pub_year <= '{year}'"
        for name in ALL_TABLES:
            if not table_has_data(self.data_dir, name):
                continue
            ref = create_table(fp, name, data_dir=dd)
            if not os.path.exists(ref.path):
                continue
            where = year_filter if name == "citation" and year_filter else ""
            self._conn.execute(
                f"CREATE OR REPLACE VIEW {name} AS "
                f"SELECT * FROM read_parquet('{ref.path}'){where}")
        if has_all_tables(self.data_dir):
            write_build_date(self.data_dir)
        self._views = True

    def execute(self, sql: str, params=None):
        self._ensure_views()
        return self._conn.execute(sql, params) if params else self._conn.execute(sql)

    def df(self, sql: str, params=None):
        return self.execute(sql, params).df()

    def fetchall(self, sql: str, params=None):
        return self.execute(sql, params).fetchall()

    def fetchone(self, sql: str, params=None):
        return self.execute(sql, params).fetchone()

    def tables(self) -> list[str]:
        self._ensure_views()
        return [r[0] for r in self._conn.execute("SHOW TABLES").fetchall()]


def main():
    """`uv run pmq`: build + compact + serve by default."""
    import argparse
    p = argparse.ArgumentParser(
        prog='pmq',
        description='PubMed database builder. Builds, compacts, and serves by default.')
    sub = p.add_subparsers(dest='command')

    q = sub.add_parser('query', help='Run a SQL query against PubMed')
    q.add_argument('sql', help='SQL query string')
    q.add_argument('--data-dir', help='Data directory')
    q.add_argument('--through', help='Date checkpoint (e.g. "2024")')

    p.add_argument('--data-dir', help='Output directory for Parquet files')
    p.add_argument('--cache-dir', help='Download cache directory')
    p.add_argument('--workers', type=int, help='Parse worker count')
    p.add_argument('--dl-jobs', type=int, help='Concurrent downloads')
    p.add_argument('--no-baseline', action='store_true')
    p.add_argument('--no-updates', action='store_true')
    p.add_argument('--no-check', action='store_true',
                    help='Skip disk space check')
    p.add_argument('--no-compact', action='store_true',
                    help='Skip compaction (keep per-worker files only)')
    p.add_argument('--disable-providing-to-community', action='store_true',
                    help='Do not serve data via P2P. Discouraged: every '
                         'builder serving makes the network faster for everyone.')

    args = p.parse_args()

    if args.command == 'query':
        with PubMed(data_dir=args.data_dir, through=args.through) as db:
            print(db.df(args.sql).to_string())
        return

    dd = Path(args.data_dir).resolve() if args.data_dir else default_data_dir()

    print(
        "\n"
        "[pmq] PubMed Database Builder\n"
        "[pmq]\n"
        "[pmq] This will download and index the complete PubMed dataset\n"
        "[pmq] (~40 million biomedical articles from NCBI). The initial\n"
        "[pmq] build takes 10-60 minutes depending on your machine and\n"
        "[pmq] network. Once built, the data is yours to query locally\n"
        "[pmq] and is shared with other researchers via P2P.\n"
        "[pmq]\n"
        f"[pmq] Data directory: {dd}\n",
        file=sys.stderr,
    )

    extra = []
    if args.cache_dir:   extra.extend(["--cache-dir", args.cache_dir])
    if args.workers:     extra.extend(["--workers", str(args.workers)])
    if args.dl_jobs:     extra.extend(["--dl-jobs", str(args.dl_jobs)])
    if args.no_baseline: extra.append("--no-baseline")
    if args.no_updates:  extra.append("--no-updates")
    if args.no_check:    extra.append("--no-check")
    ensure_build(dd, extra_args=extra or None)

    if not args.no_compact:
        do_serve = not args.disable_providing_to_community
        _compact_and_serve(dd, serve=do_serve)
        print("[pmq] Done. Data ready for queries.", file=sys.stderr)
    else:
        print("[pmq] Done (compact skipped).", file=sys.stderr)
