#!/usr/bin/env python3
"""Reproduce descriptive research tables from the public SQLite metadata core.

Usage: python3 research_brief.py --database downloaded.sqlite --output brief-output
Without --database, downloads https://psilocybin-research.com/database.php.
Uses the Python standard library. Never reads source text or runtime secrets.
"""
import argparse
import csv
from collections import Counter
import hashlib
import json
from pathlib import Path
import sqlite3
import tempfile
import urllib.request


def summarize(path):
    digest = hashlib.sha256(Path(path).read_bytes()).hexdigest()
    with sqlite3.connect(Path(path).resolve().as_uri() + '?mode=ro', uri=True) as db:
        db.row_factory = sqlite3.Row
        columns = {r['name'] for r in db.execute('PRAGMA table_info(publications)')}
        required = {'publication_year', 'publication_status', 'source_name', 'study_type', 'doi', 'pubmed_id'}
        if not required <= columns:
            raise ValueError('Input lacks required public metadata columns')
        # Runtime copies, if deliberately supplied, must exclude withheld rows.
        visible = ' WHERE hidden = 0 AND false_positive = 0' if {'hidden', 'false_positive'} <= columns else ''
        total = db.execute('SELECT COUNT(*) FROM publications' + visible).fetchone()[0]
        unknown = db.execute('SELECT COUNT(*) FROM publications' + visible + (' AND' if visible else ' WHERE') + ' publication_year IS NULL').fetchone()[0]
        years = [dict(r) for r in db.execute('SELECT publication_year AS year, COALESCE(NULLIF(TRIM(publication_status), \'\'), \'Unclassified\') AS status, COUNT(*) AS records FROM publications' + visible + ' GROUP BY publication_year, status ORDER BY publication_year, status')]
        groups = {}
        for name, column in [('statuses', 'publication_status'), ('sources', 'source_name'), ('study_types', 'study_type')]:
            groups[name] = [dict(r) for r in db.execute('SELECT COALESCE(NULLIF(TRIM(' + column + '), \'\'), \'Unclassified\') AS label, COUNT(*) AS records FROM publications' + visible + ' GROUP BY label ORDER BY records DESC, label')]
        identifiers = {}
        for key, column in [('doi', 'doi'), ('pubmed', 'pubmed_id')]:
            identifiers[key] = db.execute('SELECT COUNT(*) FROM publications' + visible + (' AND' if visible else ' WHERE') + ' ' + column + ' IS NOT NULL AND TRIM(' + column + ') != \'\'').fetchone()[0]
        return {'snapshot_sha256': digest, 'total': total, 'unknown_year': unknown, 'identifiers': identifiers, 'years': years, **groups}


def write_tables(summary, output, records):
    output.mkdir(parents=True, exist_ok=True)
    with (output / 'records.csv').open('w', encoding='utf-8', newline='') as stream:
        fields = ['id', 'publication_year', 'publication_status', 'source_name', 'study_type', 'doi', 'pubmed_id']
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        writer.writerows(records)
    summary['record_metadata_sha256'] = hashlib.sha256((output / 'records.csv').read_bytes()).hexdigest()
    (output / 'summary.json').write_text(json.dumps(summary, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')
    for name, fields in [('years', ['year', 'status', 'records']), ('statuses', ['label', 'records']), ('sources', ['label', 'records']), ('study_types', ['label', 'records'])]:
        with (output / (name + '.csv')).open('w', encoding='utf-8', newline='') as stream:
            writer = csv.DictWriter(stream, fieldnames=fields)
            writer.writeheader()
            writer.writerows(summary[name])


def summarize_records(path):
    with path.open(encoding='utf-8', newline='') as stream:
        records = list(csv.DictReader(stream))
    groups = {key: Counter((r[column] or '').strip() or 'Unclassified' for r in records)
              for key, column in [('statuses', 'publication_status'), ('sources', 'source_name'), ('study_types', 'study_type')]}
    years = Counter((int(r['publication_year']) if r['publication_year'] else None, (r['publication_status'] or '').strip() or 'Unclassified') for r in records)
    return {'total': len(records), 'unknown_year': sum(not r['publication_year'] for r in records),
            'identifiers': {key: sum(bool((r[column] or '').strip()) for r in records) for key, column in [('doi', 'doi'), ('pubmed', 'pubmed_id')]},
            'years': [{'year': year, 'status': status, 'records': count} for (year, status), count in sorted(years.items(), key=lambda item: (item[0][0] or 0, item[0][1]))],
            **{key: [{'label': label, 'records': count} for label, count in sorted(counts.items(), key=lambda item: (-item[1], item[0]))] for key, counts in groups.items()}}, records


def record_metadata(path):
    with sqlite3.connect(path.resolve().as_uri() + '?mode=ro', uri=True) as db:
        db.row_factory = sqlite3.Row
        columns = {r['name'] for r in db.execute('PRAGMA table_info(publications)')}
        visible = ' WHERE hidden = 0 AND false_positive = 0' if {'hidden', 'false_positive'} <= columns else ''
        return [dict(r) for r in db.execute('SELECT id, publication_year, publication_status, source_name, study_type, doi, pubmed_id FROM publications' + visible + ' ORDER BY id')]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    inputs = parser.add_mutually_exclusive_group()
    inputs.add_argument('--database', type=Path)
    inputs.add_argument('--records', type=Path, help='Reproduce tables from the published records.csv')
    parser.add_argument('--output', type=Path, default=Path('brief-output'))
    args = parser.parse_args()
    if args.records:
        summary, records = summarize_records(args.records)
    elif args.database:
        summary = summarize(args.database)
        records = record_metadata(args.database)
    else:
        with tempfile.TemporaryDirectory(prefix='psilo-brief-') as folder:
            path = Path(folder) / 'public-core.sqlite'
            with urllib.request.urlopen('https://psilocybin-research.com/database.php', timeout=120) as response:
                path.write_bytes(response.read())
            summary = summarize(path)
            records = record_metadata(path)
    write_tables(summary, args.output, records)
    print('Records:', summary['total'], '| Input SHA-256:', summary.get('snapshot_sha256', summary['record_metadata_sha256']))
    print('Tables:', args.output.resolve())


if __name__ == '__main__':
    main()