"""
author_analytics.py
Dedicated Analytics Engine & Database System for Author-level Word Usage in Sindhi Books & Articles.

Features:
- Completely separate database file: `author_analytics.db`
- Indexes author-level n-grams (1-grams and 2-grams)
- Tracks total words, total publications, and relative percentage frequencies per author
- Query API for Top 10 authors for any given search word or phrase
- Incremental ingestion hook for new books added to the system
- Standalone CLI to rebuild or inspect author analytics
"""

import os
import sys
import re
import glob
import time
import sqlite3
import argparse
from typing import Dict, List, Any, Optional
from collections import Counter, defaultdict

# Ensure UTF-8 stdout encoding for Windows
if sys.platform == 'win32':
    try:
        sys.stdout.reconfigure(encoding='utf-8')
    except Exception:
        pass

import sindhi_nlp

def resolve_author_database():
    local_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "author_analytics.db")
    if os.path.exists(local_path):
        return local_path
    try:
        from huggingface_hub import hf_hub_download
        token = os.environ.get("HF_TOKEN")
        print("Fetching author_analytics.db from dataset repo...")
        return hf_hub_download(
            repo_id="hellosindh/sindhi-ngram-database",
            filename="author_analytics.db",
            repo_type="dataset",
            token=token,
        )
    except Exception as e:
        print(f"Error fetching Author DB from Hub: {e}")
        return local_path

DB_FILE = resolve_author_database()


def get_db(db_path: str = DB_FILE) -> sqlite3.Connection:
    """Returns a SQLite connection to author_analytics.db with WAL mode."""
    conn = sqlite3.connect(db_path, timeout=30.0)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode = WAL;")
    conn.execute("PRAGMA synchronous = NORMAL;")
    return conn


def clean_author_name(author_str: str) -> str:
    """Strips translation notes or parenthesis from author string."""
    if not author_str:
        return "نامعلوم"
    clean = re.sub(r'\(.*?\)', '', author_str).strip()
    return clean if clean else author_str.strip()


def init_author_db(db_path: str = DB_FILE, create_indexes: bool = True) -> sqlite3.Connection:
    """Initializes schema and performance indexes for author_analytics.db."""
    conn = sqlite3.connect(db_path, timeout=30.0)
    cur = conn.cursor()

    cur.execute("PRAGMA journal_mode = WAL;")
    cur.execute("PRAGMA synchronous = NORMAL;")

    # 1. Author Profiles Table
    cur.execute("""
    CREATE TABLE IF NOT EXISTS author_profiles (
        author TEXT PRIMARY KEY,
        clean_author TEXT,
        total_words INTEGER DEFAULT 0,
        total_publications INTEGER DEFAULT 0,
        doc_types TEXT DEFAULT 'book'
    );
    """)

    # 2. Author Word Counts Table
    cur.execute("""
    CREATE TABLE IF NOT EXISTS author_word_counts (
        author TEXT,
        term TEXT,
        term_norm TEXT,
        term_heh_norm TEXT,
        n_size INTEGER,
        count INTEGER
    );
    """)

    # 3. Indexed Books Tracking Table
    cur.execute("""
    CREATE TABLE IF NOT EXISTS indexed_books (
        filename TEXT PRIMARY KEY,
        author TEXT,
        file_hash TEXT,
        total_words INTEGER,
        doc_type TEXT DEFAULT 'book',
        indexed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    );
    """)

    if create_indexes:
        create_db_indexes(cur)

    conn.commit()
    return conn


def create_db_indexes(cur: sqlite3.Cursor):
    """Creates performance indexes for sub-millisecond lookups."""
    cur.execute("CREATE INDEX IF NOT EXISTS idx_awc_norm ON author_word_counts(term_norm, n_size);")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_awc_heh ON author_word_counts(term_heh_norm, n_size);")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_awc_term ON author_word_counts(term, n_size);")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_awc_auth_term ON author_word_counts(author, term, n_size);")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_author_words ON author_profiles(total_words DESC);")


def index_single_book(filepath: str, conn: Optional[sqlite3.Connection] = None, doc_type: str = "book") -> bool:
    """
    Indexes a single book/article into author_analytics.db.
    Can be called incrementally whenever a new publication is ingested or updated.
    """
    close_at_end = False
    if conn is None:
        conn = get_db()
        close_at_end = True

    try:
        filename = os.path.basename(filepath)
        book_data = sindhi_nlp.parse_book_file(filepath)
        if not book_data:
            return False

        raw_author = book_data.get("author", "").strip() or "نامعلوم"
        author = clean_author_name(raw_author)
        text = book_data.get("text", "")
        tokens = sindhi_nlp.tokenize(text, normalize=False)
        total_tokens = len(tokens)
        if total_tokens == 0:
            return False

        hasher = sindhi_nlp.hashlib.sha256() if hasattr(sindhi_nlp, 'hashlib') else __import__('hashlib').sha256()
        with open(filepath, 'rb') as f:
            while ch := f.read(65536):
                hasher.update(ch)
        file_hash = hasher.hexdigest()

        cur = conn.cursor()

        # Check if already indexed with same hash
        cur.execute("SELECT file_hash, total_words FROM indexed_books WHERE filename = ?", (filename,))
        existing = cur.fetchone()
        if existing and existing["file_hash"] == file_hash:
            return True

        try:
            conn.execute("PRAGMA journal_mode = WAL;")
            conn.execute("PRAGMA synchronous = NORMAL;")
            conn.execute("PRAGMA cache_size = -32000;")
            conn.execute("PRAGMA temp_store = MEMORY;")
        except Exception:
            pass

        # Extract 1-grams and 2-grams
        counts_1g = Counter(tokens)
        counts_2g = Counter()
        for i in range(len(tokens) - 1):
            counts_2g[tokens[i] + ' ' + tokens[i+1]] += 1

        staging_rows = []
        norm_map = {}
        for term, cnt in counts_1g.items():
            t_norm = sindhi_nlp.normalize_sindhi(term, strip_diacritics=True, normalize_heh_variants=False)
            t_heh = sindhi_nlp.normalize_sindhi(term, strip_diacritics=True, normalize_heh_variants=True)
            norm_map[term] = (t_norm, t_heh)
            staging_rows.append((term, t_norm, t_heh, 1, cnt))

        for term, cnt in counts_2g.items():
            if cnt < 2:
                continue
            words = term.split(' ', 1)
            if len(words) == 2:
                n1 = norm_map.get(words[0]) or (words[0], words[0])
                n2 = norm_map.get(words[1]) or (words[1], words[1])
                t_norm = n1[0] + ' ' + n2[0]
                t_heh = n1[1] + ' ' + n2[1]
            else:
                t_norm = sindhi_nlp.normalize_sindhi(term, strip_diacritics=True, normalize_heh_variants=False)
                t_heh = sindhi_nlp.normalize_sindhi(term, strip_diacritics=True, normalize_heh_variants=True)
            staging_rows.append((term, t_norm, t_heh, 2, cnt))

        if staging_rows:
            batch_insert = [(author, r[0], r[1], r[2], r[3], r[4]) for r in staging_rows]
            chunk_size = 20000
            for i in range(0, len(batch_insert), chunk_size):
                cur.executemany(
                    "INSERT INTO author_word_counts (author, term, term_norm, term_heh_norm, n_size, count) VALUES (?, ?, ?, ?, ?, ?)",
                    batch_insert[i:i + chunk_size]
                )
            staging_rows.clear()

        # Upsert author profile
        cur.execute("""
        INSERT INTO author_profiles (author, clean_author, total_words, total_publications, doc_types)
        VALUES (?, ?, ?, 1, ?)
        ON CONFLICT(author) DO UPDATE SET
            total_words = total_words + excluded.total_words,
            total_publications = total_publications + 1;
        """, (author, author, total_tokens, doc_type))

        # Record indexed book
        cur.execute("""
        INSERT OR REPLACE INTO indexed_books (filename, author, file_hash, total_words, doc_type)
        VALUES (?, ?, ?, ?, ?)
        """, (filename, author, file_hash, total_tokens, doc_type))

        conn.commit()
        return True
    except Exception as e:
        print(f"[author_analytics] Error indexing {filepath}: {e}")
        return False
    finally:
        if close_at_end:
            conn.close()


def remove_book(filename: str, conn: Optional[sqlite3.Connection] = None) -> bool:
    """Removes a book tracking record from author_analytics.db."""
    close_at_end = False
    if conn is None:
        conn = get_db()
        close_at_end = True
    try:
        cur = conn.cursor()
        cur.execute("DELETE FROM indexed_books WHERE filename = ?", (filename,))
        conn.commit()
        return True
    finally:
        if close_at_end:
            conn.close()


def get_top_authors_for_term(
    term: str,
    top_n: int = 10,
    normalize: bool = True,
    heh_norm: bool = True
) -> Dict[str, Any]:
    """
    Finds the top N authors who used a given word or phrase the most.
    Calculates both raw occurrences and relative frequency percentage.
    """
    clean_term = term.strip()
    if not clean_term:
        return {"term": term, "top_authors": []}

    words = clean_term.split()
    n_size = len(words)
    if n_size > 2:
        n_size = 2  # Support up to bigrams

    conn = get_db()
    cur = conn.cursor()

    if heh_norm:
        search_key = sindhi_nlp.normalize_sindhi(clean_term, strip_diacritics=True, normalize_heh_variants=True)
        sql = """
        SELECT a.author, SUM(a.count) as total_occurrences, p.total_words, p.total_publications
        FROM author_word_counts a
        LEFT JOIN author_profiles p ON a.author = p.author
        WHERE a.term_heh_norm = ? AND a.n_size = ?
        GROUP BY a.author
        ORDER BY total_occurrences DESC
        LIMIT ?
        """
    elif normalize:
        search_key = sindhi_nlp.normalize_sindhi(clean_term, strip_diacritics=True, normalize_heh_variants=False)
        sql = """
        SELECT a.author, SUM(a.count) as total_occurrences, p.total_words, p.total_publications
        FROM author_word_counts a
        LEFT JOIN author_profiles p ON a.author = p.author
        WHERE a.term_norm = ? AND a.n_size = ?
        GROUP BY a.author
        ORDER BY total_occurrences DESC
        LIMIT ?
        """
    else:
        search_key = clean_term
        sql = """
        SELECT a.author, SUM(a.count) as total_occurrences, p.total_words, p.total_publications
        FROM author_word_counts a
        LEFT JOIN author_profiles p ON a.author = p.author
        WHERE a.term = ? AND a.n_size = ?
        GROUP BY a.author
        ORDER BY total_occurrences DESC
        LIMIT ?
        """

    cur.execute(sql, (search_key, n_size, top_n))
    rows = cur.fetchall()
    conn.close()

    top_authors = []
    for rank, r in enumerate(rows, 1):
        author_name = r["author"] or "نامعلوم"
        occ = r["total_occurrences"] or 0
        tot_words = r["total_words"] or 1
        pct = (occ / tot_words) * 100.0

        top_authors.append({
            "rank": rank,
            "author": author_name,
            "count": occ,
            "total_words": tot_words,
            "percentage": round(pct, 5),
            "percentage_str": f"{pct:.4f}%",
            "total_publications": r["total_publications"] or 1
        })

    return {
        "term": term,
        "n_size": n_size,
        "search_key": search_key,
        "top_authors": top_authors
    }


def build_full_database(scan_dirs: List[str] = None, db_path: str = DB_FILE, batch_commit_size: int = 100000) -> bool:
    """
    Fast batch processor that builds author_analytics.db from all book and article files in `books/`.
    Optimized:
    - Pre-aggregates in memory (takes ~25s for 9,210 files)
    - Keeps all 1-grams (count >= 1) and all meaningful 2-grams (count >= 2)
    - Bulk inserts raw rows without active indexes in PRAGMA synchronous = OFF mode
    - Creates indexes after insert
    - Total build time: ~45-60 seconds!
    """
    if scan_dirs is None:
        scan_dirs = ["books"]

    print(f"[author_analytics] Initializing database at: {db_path}")
    for ext in ["", "-wal", "-shm"]:
        target = db_path + ext
        if os.path.exists(target):
            try:
                os.remove(target)
            except Exception:
                pass

    conn = init_author_db(db_path, create_indexes=False)
    cur = conn.cursor()

    cur.execute("PRAGMA synchronous = OFF;")
    cur.execute("PRAGMA journal_mode = MEMORY;")
    cur.execute("PRAGMA cache_size = -128000;")  # 128MB cache

    files = []
    for d in scan_dirs:
        if os.path.exists(d):
            files.extend(glob.glob(os.path.join(d, "*.txt")))

    files = sorted(list(set(files)))
    total_files = len(files)
    print(f"[author_analytics] Found {total_files:,} files to process.")

    t0 = time.time()
    author_1g = defaultdict(Counter)
    author_2g = defaultdict(Counter)
    author_totals = Counter()
    author_pubs = Counter()
    author_doc_types = defaultdict(set)
    indexed_records = []

    for idx, f in enumerate(files, 1):
        data = sindhi_nlp.parse_book_file(f)
        if not data:
            continue

        raw_author = data.get("author", "").strip() or "نامعلوم"
        author = clean_author_name(raw_author)
        text = data.get("text", "")
        tokens = sindhi_nlp.tokenize(text, normalize=False)
        num_tokens = len(tokens)
        if num_tokens == 0:
            continue

        filename = os.path.basename(f)
        doc_type = "article" if filename.startswith("article_") else "book"

        author_totals[author] += num_tokens
        author_pubs[author] += 1
        author_doc_types[author].add(doc_type)

        hasher = sindhi_nlp.hashlib.sha256() if hasattr(sindhi_nlp, 'hashlib') else __import__('hashlib').sha256()
        with open(f, 'rb') as fp:
            while ch := fp.read(65536):
                hasher.update(ch)
        f_hash = hasher.hexdigest()
        indexed_records.append((filename, author, f_hash, num_tokens, doc_type))

        for tok in tokens:
            author_1g[author][tok] += 1
        for i in range(len(tokens) - 1):
            author_2g[author][tokens[i] + ' ' + tokens[i+1]] += 1

        if idx % 2000 == 0 or idx == total_files:
            print(f"  Processed {idx:,}/{total_files:,} files ({time.time()-t0:.1f}s)...")

    scan_time = time.time() - t0
    print(f"[author_analytics] Scanning complete in {scan_time:.1f}s. Preparing bulk insert...")

    # 1. Insert author profiles
    profiles_data = [
        (author, author, author_totals[author], author_pubs[author], ",".join(author_doc_types[author]))
        for author in author_totals
    ]
    cur.executemany("""
    INSERT INTO author_profiles (author, clean_author, total_words, total_publications, doc_types)
    VALUES (?, ?, ?, ?, ?)
    """, profiles_data)

    # 2. Insert indexed books
    cur.executemany("""
    INSERT INTO indexed_books (filename, author, file_hash, total_words, doc_type)
    VALUES (?, ?, ?, ?, ?)
    """, indexed_records)

    # 3. Bulk insert 1-grams (all words)
    print("[author_analytics] Inserting 1-grams...")
    insert_sql = """
    INSERT INTO author_word_counts (author, term, term_norm, term_heh_norm, n_size, count)
    VALUES (?, ?, ?, ?, ?, ?)
    """
    batch = []
    total_1g_rows = 0
    for author, counter in author_1g.items():
        for term, cnt in counter.items():
            t_norm = sindhi_nlp.normalize_sindhi(term, strip_diacritics=True, normalize_heh_variants=False)
            t_heh = sindhi_nlp.normalize_sindhi(term, strip_diacritics=True, normalize_heh_variants=True)
            batch.append((author, term, t_norm, t_heh, 1, cnt))
            if len(batch) >= batch_commit_size:
                cur.executemany(insert_sql, batch)
                total_1g_rows += len(batch)
                batch = []
    if batch:
        cur.executemany(insert_sql, batch)
        total_1g_rows += len(batch)
    print(f"[author_analytics] Inserted {total_1g_rows:,} unigram entries.")

    # 4. Bulk insert 2-grams (phrases with count >= 2 for meaningful author idioms)
    print("[author_analytics] Inserting 2-grams (frequency >= 2)...")
    batch = []
    total_2g_rows = 0
    for author, counter in author_2g.items():
        for term, cnt in counter.items():
            if cnt >= 2:
                t_norm = sindhi_nlp.normalize_sindhi(term, strip_diacritics=True, normalize_heh_variants=False)
                t_heh = sindhi_nlp.normalize_sindhi(term, strip_diacritics=True, normalize_heh_variants=True)
                batch.append((author, term, t_norm, t_heh, 2, cnt))
                if len(batch) >= batch_commit_size:
                    cur.executemany(insert_sql, batch)
                    total_2g_rows += len(batch)
                    batch = []
    if batch:
        cur.executemany(insert_sql, batch)
        total_2g_rows += len(batch)
    print(f"[author_analytics] Inserted {total_2g_rows:,} bigram entries.")

    conn.commit()

    # 5. Create indexes after data insertion
    print("[author_analytics] Creating performance indexes...")
    t_idx = time.time()
    create_db_indexes(cur)
    conn.commit()
    print(f"[author_analytics] Indexes created in {time.time()-t_idx:.1f}s.")

    # Switch to WAL mode for concurrent readers and writer
    cur.execute("PRAGMA journal_mode = WAL;")
    cur.execute("PRAGMA synchronous = NORMAL;")
    conn.commit()
    conn.close()

    total_time = time.time() - t0
    print(f"[author_analytics] Successfully built {db_path} ({total_1g_rows + total_2g_rows:,} total entries) in {total_time:.1f}s!")
    return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Author Analytics Engine for Sindhi Literature")
    parser.add_argument("--build", action="store_true", help="Rebuild the author analytics database from books/")
    parser.add_argument("--query", type=str, help="Query top 10 authors for a Sindhi word or phrase")
    parser.add_argument("--exact", action="store_true", help="Exact matching (disable normalization)")
    args = parser.parse_args()

    if args.build:
        build_full_database()
    elif args.query:
        norm = not args.exact
        res = get_top_authors_for_term(args.query, top_n=10, normalize=norm, heh_norm=norm)
        print(f"\nTop 10 Authors for '{args.query}':")
        for a in res["top_authors"]:
            print(f"  #{a['rank']} {a['author']}: {a['count']:,} occurrences ({a['percentage_str']}, total words: {a['total_words']:,})")
    else:
        parser.print_help()