#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Export selected tables from an SQLite DB to JSON and extract media files.
Usage: python export_sqlite.py [path/to/db] [output_dir]
Defaults: db -> ZYAKUKGN.db (current dir), output_dir -> ./export

Behavior:
- Exports tables `info`, `content`, `list` (if present) to JSON files in the output dir.
  BLOB fields in those JSON files are encoded as base64 strings.
- Extracts rows from `media` table and writes files into output_dir/media.
  File name uses the `name` field. Extension mapping: 1->jpg, 2->png, 4->svg. Other types -> .bin
  The file contents are taken from the `main` field. If `main` is a BLOB it's written directly.
  If `main` is text it will try to base64-decode it; if that fails it's written as UTF-8 text
  (useful for SVG stored as text).
"""

import sqlite3
import json
import os
import sys
import base64
import re
from pathlib import Path

EXT_MAP = {1: 'jpg', 2: 'png', 4: 'svg'}


def safe_filename(name: str) -> str:
    if not name:
        return 'unnamed'
    # remove path chars and keep reasonable subset
    name = str(name)
    name = re.sub(r'[\\/:*?"<>|]+', '_', name)
    name = name.strip()
    if not name:
        return 'unnamed'
    return name


def table_exists(conn, name: str) -> bool:
    c = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (name,))
    return c.fetchone() is not None


def rows_to_serializable(cols, rows):
    items = []
    for r in rows:
        d = {}
        for col, val in zip(cols, r):
            if isinstance(val, (bytes, bytearray)):
                d[col] = base64.b64encode(bytes(val)).decode('ascii')
            else:
                d[col] = val
        items.append(d)
    return items


def export_table_json(conn, table, out_dir):
    if not table_exists(conn, table):
        print(f"Table '{table}' not found, skipping.")
        return
    cur = conn.execute(f"SELECT * FROM {table}")
    cols = [d[0] for d in cur.description]
    rows = cur.fetchall()
    data = rows_to_serializable(cols, rows)
    out_file = out_dir / f"{table}.json"
    with out_file.open('w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    print(f"Wrote {len(data)} rows to {out_file}")


def extract_media(conn, out_dir):
    table = 'media'
    if not table_exists(conn, table):
        print("Table 'media' not found, skipping media extraction.")
        return
    media_dir = out_dir / 'media'
    media_dir.mkdir(parents=True, exist_ok=True)

    cur = conn.execute(f"SELECT * FROM {table}")
    cols = [d[0] for d in cur.description]
    rows = cur.fetchall()

    used = {}
    wrote = 0
    for r in rows:
        row = dict(zip(cols, r))
        name = row.get('name') or row.get('file_name') or row.get('title') or 'unnamed'
        typ = row.get('type')
        main = row.get('main')

        ext = EXT_MAP.get(typ, 'bin')
        base = safe_filename(name)
        filename = f"{base}.{ext}"
        path = media_dir / filename
        # avoid overwriting
        i = 1
        while path.exists():
            path = media_dir / f"{base}_{i}.{ext}"
            i += 1

        try:
            if main is None:
                print(f"Skipping media row with name={name}: no main data")
                continue

            if isinstance(main, (bytes, bytearray)):
                # bytes: write directly. For svg try to write as text if possible
                if ext == 'svg':
                    try:
                        text = main.decode('utf-8')
                        path.write_text(text, encoding='utf-8')
                    except Exception:
                        path.write_bytes(bytes(main))
                else:
                    path.write_bytes(bytes(main))

            elif isinstance(main, str):
                s = main.strip()
                # svg xml text
                if ext == 'svg' or s.startswith('<'):
                    path.write_text(s, encoding='utf-8')
                else:
                    # try base64 decode
                    try:
                        b = base64.b64decode(s, validate=True)
                        path.write_bytes(b)
                    except Exception:
                        # fallback: write raw utf-8
                        path.write_text(s, encoding='utf-8')

            else:
                # fallback: convert to string
                path.write_text(str(main), encoding='utf-8')

            used[filename] = {'name': name, 'type': typ}
            wrote += 1
        except Exception as e:
            print(f"Failed writing media {name}: {e}")

    # write index for media
    idx_file = media_dir / 'media_index.json'
    with idx_file.open('w', encoding='utf-8') as f:
        json.dump(used, f, ensure_ascii=False, indent=2)
    print(f"Extracted {wrote} media files to {media_dir}")


def main():
    db_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path('data.db')
    out_dir = Path(sys.argv[2]) if len(sys.argv) > 2 else Path('export')

    if not db_path.exists():
        print(f"Database file not found: {db_path}")
        sys.exit(2)

    out_dir.mkdir(parents=True, exist_ok=True)

    conn = sqlite3.connect(str(db_path))
    try:
        for t in ['info', 'content', 'list', 'search']:
            export_table_json(conn, t, out_dir)
        extract_media(conn, out_dir)
    finally:
        conn.close()


if __name__ == '__main__':
    main()
