import time
import json
import requests
import pandas as pd
from sqlalchemy import create_engine, text
from huggingface_hub import HfFileSystem

DB_URI = "postgresql://postgres:password@localhost:5432/postgres"
engine = create_engine(DB_URI)
fs = HfFileSystem()

import io
import csv

session = requests.Session()
adapter = requests.adapters.HTTPAdapter(pool_connections=20, pool_maxsize=20)
session.mount('https://', adapter)

def copy_to_db(table_name, df, engine):
    """Fast bulk insert using PostgreSQL COPY via psycopg2."""
    if df.empty:
        return
        
    # Clean data to prevent Postgres COPY errors (carriage returns, null bytes, and float64 floats)
    for col in df.columns:
        if str(df[col].dtype) == 'float64':
            df[col] = df[col].astype('Int64')
        elif str(df[col].dtype) == 'object':
            df[col] = df[col].apply(lambda x: x.replace('\r', ' ').replace('\0', '') if isinstance(x, str) else x)

    # Convert string JSON to ensure no weird quoting issues
    buffer = io.StringIO()
    df.to_csv(buffer, index=False, header=False, sep='\t', quoting=csv.QUOTE_MINIMAL, na_rep='\\N')
    buffer.seek(0)
    
    # We use psycopg2 raw connection
    conn = engine.raw_connection()
    try:
        with conn.cursor() as cur:
            columns = ', '.join(df.columns)
            sql = f"COPY {table_name} ({columns}) FROM STDIN WITH (FORMAT CSV, DELIMITER '\t', NULL '\\N', QUOTE '\"', ESCAPE '\"')"
            cur.copy_expert(sql, buffer)
        conn.commit()
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        conn.close()

def build_path_mapping(fs):
    print("Building path mapping from HuggingFace (this might take a few seconds)...")
    paths = fs.glob("datasets/AuthenticIlm/Shamela4_Full_DB/*/*")
    mapping = {}
    for p in paths:
        parts = p.split('/')
        if len(parts) >= 4:
            book_dir = parts[-1]
            try:
                book_id_str = book_dir.split('__')[0]
                if book_id_str.isdigit():
                    mapping[int(book_id_str)] = f"{parts[-2]}/{parts[-1]}/pages.jsonl"
            except:
                pass
    print(f"Path mapping built for {len(mapping)} books.")
    return mapping

def process_book(book_id, file_path):
    if not file_path:
        return False, "Not found in mapping"
        
    # TOC insertion moved to the end of process_book to prevent FK violations
    
    url = f"https://huggingface.co/datasets/AuthenticIlm/Shamela4_Full_DB/resolve/main/{file_path}"
    try:
        with session.get(url, stream=True, timeout=30) as response:
            if response.status_code != 200:
                return False, "Download failed"
        
            pages = []
            
            for line in response.iter_lines(chunk_size=65536):
                if line:
                    try:
                        page_data = json.loads(line)
                        if 'services_raw' in page_data and page_data['services_raw'] is not None:
                            page_data['services_raw'] = json.dumps(page_data['services_raw'])
                        
                        pages.append(page_data)
                        
                        if len(pages) >= 5000:
                            pages_df = pd.DataFrame(pages)
                            cols = ['page_id', 'book_id', 'shamela_page_id', 'part', 'page_num', 'sequence_num', 'body', 'footnotes', 'hints', 'services_raw']
                            pages_df = pages_df[[c for c in cols if c in pages_df.columns]]
                            copy_to_db('pages', pages_df, engine)
                            pages = []
                    except Exception as e:
                        pass
                        
            if len(pages) > 0:
                try:
                    pages_df = pd.DataFrame(pages)
                    cols = ['page_id', 'book_id', 'shamela_page_id', 'part', 'page_num', 'sequence_num', 'body', 'footnotes', 'hints', 'services_raw']
                    pages_df = pages_df[[c for c in cols if c in pages_df.columns]]
                    copy_to_db('pages', pages_df, engine)
                except Exception as e:
                    print(f"Error saving final pages for book {book_id}: {e}")

            # Now that pages are inserted, we can safely insert TOC which depends on pages(page_id)
            toc_file_path = file_path.replace("pages.jsonl", "toc.jsonl")
            toc_url = f"https://huggingface.co/datasets/AuthenticIlm/Shamela4_Full_DB/resolve/main/{toc_file_path}"
            toc_entries = []
            try:
                with session.get(toc_url, stream=True, timeout=30) as toc_response:
                    if toc_response.status_code == 200:
                        for line in toc_response.iter_lines(chunk_size=65536):
                            if line:
                                try:
                                    toc_entries.append(json.loads(line))
                                except Exception as e:
                                    pass
            except Exception as e:
                print(f"Error downloading TOC for {book_id}: {e}")
                
            if len(toc_entries) > 0:
                try:
                    # Ensure correct columns mapping if json has extra fields
                    toc_df = pd.DataFrame(toc_entries)
                    expected_cols = ['title_id', 'book_id', 'page_id', 'parent_id', 'shamela_title_id', 'title_text']
                    toc_df = toc_df[[c for c in expected_cols if c in toc_df.columns]]
                    
                    # Fix Shamela edge-case where page_id=0 for general headings, which violates FK constraints
                    if 'page_id' in toc_df.columns:
                        toc_df.loc[toc_df['page_id'] == 0, 'page_id'] = None
                        
                    copy_to_db('toc', toc_df, engine)
                except Exception as e:
                    print(f"Error saving TOC for book {book_id}: {e}")
                
            return True, "Success"
    except Exception as e:
        return False, str(e)

def process_wrapper(args):
    book_id, file_path = args
    print(f"Processing book {book_id}...")
    success, msg = process_book(book_id, file_path)
    if success:
        print(f"Book {book_id} processed successfully.")
    else:
        print(f"Failed to process book {book_id}: {msg}")

def load_metadata():
    import numpy as np
    print("Downloading Categories...")
    df = pd.read_parquet("https://huggingface.co/datasets/AuthenticIlm/Shamela4_Full_DB/resolve/main/_meta/categories.parquet")
    if 'id' in df.columns: df.rename(columns={'id': 'category_id'}, inplace=True)
    if 'name_ar' in df.columns and 'category_name_ar' not in df.columns: df.rename(columns={'name_ar': 'category_name_ar'}, inplace=True)
    if 'category_name' in df.columns and 'category_name_ar' not in df.columns: df.rename(columns={'category_name': 'category_name_ar'}, inplace=True)
    cols = ['category_id', 'category_name_ar', 'name_en', 'sort_order']
    df[[c for c in cols if c in df.columns]].to_sql("categories", engine, if_exists="append", index=False)
    
    print("Downloading Authors...")
    df = pd.read_parquet("https://huggingface.co/datasets/AuthenticIlm/Shamela4_Full_DB/resolve/main/_meta/authors.parquet")
    if 'id' in df.columns: df.rename(columns={'id': 'author_id'}, inplace=True)
    cols = ['author_id', 'name_ar', 'death_hijri', 'death_hijri_text', 'alpha_sort', 'biography']
    df[[c for c in cols if c in df.columns]].to_sql("authors", engine, if_exists="append", index=False)
    
    print("Downloading Books Metadata...")
    df = pd.read_parquet("https://huggingface.co/datasets/AuthenticIlm/Shamela4_Full_DB/resolve/main/_meta/book_metadata.parquet")
    cols = ['book_id', 'shamela_id', 'title_ar', 'book_type', 'book_type_label', 'category_id', 'main_author_id', 'main_author_name_ar', 'main_author_death_hijri', 'main_author_death_hijri_text', 'authors_text', 'hijri_era', 'printed', 'is_hidden', 'parent_id', 'group_id', 'version_major', 'version_minor', 'betaka_text', 'meta', 'volume_count_observed', 'has_multi_part', 'authors_json']
    df = df[[c for c in cols if c in df.columns]]
    if 'meta' in df.columns:
        df['meta'] = df['meta'].apply(lambda x: json.dumps(x.tolist()) if isinstance(x, np.ndarray) else (json.dumps(x) if isinstance(x, dict) else (x if pd.isna(x) else json.dumps(x))))
    if 'authors_json' in df.columns:
        df['authors_json'] = df['authors_json'].apply(lambda x: json.dumps(x.tolist()) if isinstance(x, np.ndarray) else (json.dumps(x) if isinstance(x, dict) else (x if pd.isna(x) else json.dumps(x))))
    df.to_sql("books", engine, if_exists="append", index=False)
    print("Metadata loading complete!")

def run_worker():
    print("Ingest Worker Started...")
    from concurrent.futures import ThreadPoolExecutor
    path_mapping = None
    
    while True:
        try:
            # Get list of all books
            df_books = pd.read_sql("SELECT book_id FROM books", engine)
            if df_books.empty:
                print("No books found in DB. Auto-loading metadata...")
                try:
                    load_metadata()
                except Exception as e:
                    print(f"Failed to load metadata: {e}")
                    time.sleep(60)
                continue
                
            if path_mapping is None:
                path_mapping = build_path_mapping(fs)
                
            all_books = set(df_books['book_id'].tolist())
            
            # Get list of processed books
            df_processed = pd.read_sql("SELECT DISTINCT book_id FROM pages", engine)
            processed_books = set(df_processed['book_id'].tolist())
            
            pending_books = list(all_books - processed_books)
            print(f"Total books: {len(all_books)}, Processed: {len(processed_books)}, Pending: {len(pending_books)}")
            
            if not pending_books:
                print("All books processed. Creating full-text index if it doesn't exist...")
                with engine.begin() as conn:
                    conn.execute(text("CREATE INDEX IF NOT EXISTS pages_body_idx ON pages USING GIN (to_tsvector('arabic', body));"))
                print("Index created. Sleeping for 1 hour...")
                time.sleep(3600)
                continue
                
            # If we have pending books, drop the index for faster ingestion
            with engine.begin() as conn:
                conn.execute(text("DROP INDEX IF EXISTS pages_body_idx;"))
                
            # Prepare arguments for multiprocessing
            tasks = [(b, path_mapping.get(b)) for b in pending_books]
            
            # Process books concurrently with ThreadPoolExecutor
            with ThreadPoolExecutor(max_workers=5) as executor:
                executor.map(process_wrapper, tasks)
                
            print("Batch finished. Checking again...")
            time.sleep(5)
                
        except Exception as e:
            print(f"Worker Error: {e}")
            time.sleep(60)

if __name__ == "__main__":
    time.sleep(10)
    run_worker()