#!/usr/bin/env python3
"""Rebuild all download files (PDF/EPUB/MOBI) from D1 and upload to R2.

Usage: python3 rebuild_downloads.py [--dry-run] [--langs en,es,pt,hi]
"""
import json, subprocess, os, tempfile, shutil, argparse, hashlib, gzip, io
from pathlib import Path
from collections import defaultdict

# ── config ──
LANGUAGES = ['en','es','pt','hi']
BOOKS = {
    'bg': ('Bhagavad-gita', 'bg'), 'iso': ('Sri Isopanisad', 'iso'),
    'noi': ('Nectar of Instruction', 'noi'), 'nod': ('Nectar of Devotion', 'nod'),
    'tlc': ('Teachings of Lord Chaitanya', 'tlc'), 'kb': ('Krsna Book', 'kb'),
    'tqk': ('Teachings of Queen Kunti', 'tqk'), 'ssr': ('Science of Self-Realization', 'ssr'),
    'pop': ('Path of Perfection', 'pop'), 'krp': ('Krsna Reservoir of Pleasure', 'krp'),
    'ttp': ('Transcendental Teachings of Prahlad', 'ttp'),
    'poy': ('Perfection of Yoga', 'poy'), 'bbd': ('Beyond Birth and Death', 'bbd'),
    'ej': ('Easy Journey', 'ej'), 'ekc': ('Elevation to Krsna Consciousness', 'ekc'),
    'owk': ('On the Way to Krsna', 'owk'), 'pqpa': ('Perfect Questions Perfect Answers', 'pqpa'),
    'rv': ('Raja-vidya', 'rv'), 'lob': ('Light of the Bhagavata', 'lob'),
    'lcfl': ('Lord Chaitanya in Five Features', 'lcfl'),
}
SB_CANTOS = list(range(1,11))
CC_LILAS = ['adi','madhya','antya']

R2_BUCKET = 'vedabase-files'
OUT_DIR = Path('/tmp/vedabase_downloads')

# ── D1 helpers ──
def d1_query(sql):
    r = subprocess.run(['npx','wrangler','d1','execute','vedabase-search-db','--remote','--json','--command',sql],
                       capture_output=True, text=True, timeout=60)
    return json.loads(r.stdout)[0]['results']

def fetch_verses(book, lang, extra_clause=''):
    where = f"book='{book}' AND lang='{lang}'"
    if extra_clause: where += f' AND {extra_clause}'
    return d1_query(f"SELECT ref, verse_text, translation, purport FROM verses WHERE {where} ORDER BY ref")

# ── Markdown generation ──
def verses_to_md(title, verses):
    md = f'# {title}\n\n'
    for v in verses:
        ref = v['ref']
        ch = ref.split('/')[1] if '/' in ref else ''
        md += f'## {ref}\n\n'
        if v.get('translation'):
            md += f"{v['translation']}\n\n"
        if v.get('purport'):
            md += f"{v['purport']}\n\n"
        md += '\n'
    return md

# ── EPUB generation ──
def md_to_epub(title, md_text, out_path):
    import ebooklib.epub as epub
    book = epub.EpubBook()
    book.set_identifier(str(hash(title)))
    book.set_title(title)
    book.set_language('en')
    book.add_author('A.C. Bhaktivedanta Swami Prabhupada')
    
    chapter = epub.EpubHtml(title=title, file_name='content.xhtml', lang='en')
    chapter.content = f'<html><body><h1>{title}</h1>{md_to_html(md_text)}</body></html>'
    book.add_item(chapter)
    book.spine = ['nav', chapter]
    book.toc = [epub.Link('content.xhtml', title, 'main')]
    book.add_item(epub.EpubNcx())
    book.add_item(epub.EpubNav())
    
    epub.write_epub(str(out_path), book)

def md_to_html(md):
    """Basic markdown to HTML."""
    import re
    html = md
    html = re.sub(r'^# (.+)$', r'<h1>\1</h1>', html, flags=re.M)
    html = re.sub(r'^## (.+)$', r'<h2>\1</h2>', html, flags=re.M)
    html = re.sub(r'\*\*(.+?)\*\*', r'<b>\1</b>', html)
    html = html.replace('\n\n', '</p><p>')
    html = f'<p>{html}</p>'
    html = html.replace('<p></p>', '')
    html = re.sub(r'<p>(<h[12]>.*?</h[12]>)</p>', r'\1', html)
    return html

# ── PDF generation ──
def md_to_pdf(md_text, out_path):
    with tempfile.NamedTemporaryFile(suffix='.md', mode='w', delete=False) as f:
        f.write(md_text)
    r = subprocess.run(['pandoc', f.name, '-o', str(out_path), '--pdf-engine=xelatex',
                    '-V', 'geometry:margin=2cm'],
                   capture_output=True, timeout=120)
    os.unlink(f.name)
    # pandoc exits 43 on font warnings but PDF is generated
    return out_path.exists()

# ── MOBI generation ──
def epub_to_mobi(epub_path, out_path):
    subprocess.run(['ebook-convert', str(epub_path), str(out_path)],
                   capture_output=True, timeout=120)

# ── R2 upload ──
def upload_r2(key, filepath):
    r = subprocess.run(['npx','wrangler','r2','object','put',f'{R2_BUCKET}/{key}',
                        '--file', str(filepath), '--remote'], capture_output=True, text=True, timeout=120)
    return r.returncode == 0

# ── Main ──
def build_book(book_id, book_title, verses, lang):
    """Build PDF+EPUB+MOBI for a book."""
    base = OUT_DIR / lang
    base.mkdir(parents=True, exist_ok=True)
    title = book_title if lang == 'en' else f'{book_title} ({lang})'
    
    md = verses_to_md(title, verses)
    stem = f'{book_id}'
    
    # EPUB
    epub_path = base / f'{stem}.epub'
    md_to_epub(title, md, epub_path)
    
    # MOBI
    mobi_path = base / f'{stem}.mobi'
    epub_to_mobi(epub_path, mobi_path)
    
    # PDF
    pdf_path = base / f'{stem}.pdf'
    md_to_pdf(md, pdf_path)
    
    return epub_path, mobi_path, pdf_path

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dry-run', action='store_true')
    parser.add_argument('--langs', type=str, default='en,es,pt,hi')
    args = parser.parse_args()
    langs = [l.strip() for l in args.langs.split(',')]
    dry = args.dry_run
    
    todo = []
    # Small books
    for book_id, (title, _) in BOOKS.items():
        for lang in langs:
            t = title if lang == 'en' else f'{title} ({lang})'
            todo.append(('book', book_id, t, book_id, lang))
    
    # SB cantos
    for canto in SB_CANTOS:
        for lang in langs:
            t = f'Srimad-Bhagavatam Canto {canto}'
            todo.append(('sb', f'sb-canto-{canto}', t, 'sb', lang,
                        f"ref LIKE 'sb/{canto}/%'"))
    
    # CC lilas
    for lila in CC_LILAS:
        for lang in langs:
            t = f'Caitanya-caritamrta {lila.title()}-lila'
            todo.append(('cc', f'cc-{lila}', t, 'cc', lang,
                        f"ref LIKE 'cc/{lila}/%'"))
    
    print(f"Build {len(todo)} items ({len(langs)} langs) {'[DRY RUN]' if dry else ''}")
    
    for item in todo:
        if len(item) == 5:
            _, file_id, title, book, lang = item
            extra = None
        else:
            _, file_id, title, book, lang, extra = item
        
        print(f"  {file_id} [{lang}]...", end=' ', flush=True)
        
        if dry:
            print("DRY RUN")
            continue
        
        verses = fetch_verses(book, lang, extra)
        if not verses:
            print(f"0 verses, SKIP")
            continue
        
        ep, mb, pdf = build_book(file_id, title, verses, lang)
        
        for ext, path in [('epub',ep), ('mobi',mb), ('pdf',pdf)]:
            key = f'downloads/{file_id}.{ext}'
            ok = upload_r2(key, path)
            if not ok:
                print(f"UPLOAD FAIL {key}")
        
        print(f"OK ({len(verses)} verses)")

if __name__ == '__main__':
    main()