"""
HSEQ Kennisbank Indexer
PDF indexing pipeline met PyMuPDF + SQLite FTS5
"""

import os
import re
import sqlite3
import fitz  # PyMuPDF
from datetime import datetime

DB_PATH = os.path.join(os.path.dirname(__file__), 'hseq_kennisbank.db')

# Categorie mapping op basis van filename patterns
CATEGORY_RULES = [
    (r'^AI-\d+', 'Arbowetgeving'),
    (r'hitte|warmte|heat', 'Hitte & Warmte'),
    (r'asbest', 'Asbest'),
    (r'chroom|cr[- ]?6', 'Chroom-6'),
    (r'stof|dust', 'Inhaleerbaar Stof'),
    (r'gevaarlijk|chemical|chemische', 'Gevaarlijke Stoffen'),
    (r'beeldscherm', 'Beeldschermwerk'),
    (r'geluid|lawaai|trilling', 'Geluid & Trillingen'),
    (r'elektrisch|elektromagnetisch', 'Elektrische Veiligheid'),
    (r'fysiek|belasting|tillen|dragen', 'Fysieke Belasting'),
    (r'biologisch|agentia|bacterie|virus', 'Biologische Agentia'),
    (r'pandemie|covid|corona', 'Pandemie & Infectieziekten'),
    (r'brand|blus|ontruiming|nood', 'Brand & Ontruiming'),
    (r'machine|gereedschap|werktuig', 'Machines & Gereedschap'),
    (r'val|hoogte|dalig|veiligheidsladder', 'Valgevaar & Hoogwerk'),
    (r'besloten|ruimte|confined', 'Besloten Ruimten'),
    (r'explosief|atmosfeer|ex', 'Explosieve Atmosfeer (ATEX)'),
    (r'straling|uv|laser', 'Straling'),
    (r'klimaat|binnenmilieu|ventilatie', 'Klimaat & Binnenmilieu'),
    (r'psychisch|stress|burnout|pesten', 'Psychosociale Arbeidsbelasting'),
    (r'risico|analyse|beoordeling', 'Risicobeoordeling'),
    (r'PBM|persoonlijke|bescherming', 'Persoonlijke Beschermingsmiddelen'),
    (r'alcohol|drugs|medicijn', 'Alcohol, Drugs & Medicijnen'),
    (r'rusttijd|arbeidstijd|overwerk', 'Arbeidstijden'),
    (r'continuïteit|crisis|noodorg', 'Bedrijfscontinuïteit & Crisis'),
    (r'training|opleiding|instructie|voorlichting', 'Training & Opleiding'),
    (r'preventie|veiligheidscultuur|gedrag', 'Veiligheidscultuur'),
    (r'inspectie|audit|toezicht|handhaving', 'Inspectie & Handhaving'),
    (r'ergonomie|werkplek|balie|kassa', 'Ergonomie & Werkplek'),
    (r'bestrijdingsmiddel|gewasbescherming', 'Bestrijdingsmiddelen'),
    (r'wanden|vloer|openingen|vallen', 'Wand- & Vloeropeningen'),
    (r'communicatie|risico', 'Communicatie & Risico'),
    (r'ziekte|beroepsziekte', 'Beroepsziekten'),
    (r'bedrijfsruimte|gebouw', 'Bedrijfsruimten'),
    (r'handboek|handleiding|gids', 'Handboeken & Handleidingen'),
    (r'basisboek', 'Basisboeken & Referentiemateriaal'),
]

def categorize(filename):
    """Bepaal categorie op basis van filename"""
    fn_lower = filename.lower()
    for pattern, category in CATEGORY_RULES:
        if re.search(pattern, fn_lower, re.IGNORECASE):
            return category
    return 'Overig'


def get_db():
    """Get database connection"""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def init_db():
    """Initialize database schema"""
    conn = get_db()
    cursor = conn.cursor()

    cursor.executescript('''
        CREATE TABLE IF NOT EXISTS documents (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT NOT NULL,
            filepath TEXT NOT NULL,
            title TEXT,
            author TEXT,
            pages INTEGER,
            file_size INTEGER,
            created_date TEXT,
            indexed_date TEXT,
            category TEXT,
            tags TEXT,
            full_text TEXT
        );

        CREATE TABLE IF NOT EXISTS chunks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            document_id INTEGER REFERENCES documents(id),
            chunk_index INTEGER,
            content TEXT,
            char_start INTEGER DEFAULT 0,
            char_end INTEGER DEFAULT 0,
            FOREIGN KEY(document_id) REFERENCES documents(id)
        );

        CREATE TABLE IF NOT EXISTS hseq_actions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            source_title TEXT NOT NULL,
            action_item TEXT NOT NULL,
            priority TEXT DEFAULT 'HOOG',
            status TEXT DEFAULT 'Open',
            linked_article_url TEXT DEFAULT '',
            relevance_score INTEGER DEFAULT 0,
            category TEXT DEFAULT '',
            created_at TEXT DEFAULT CURRENT_TIMESTAMP,
            closed_at TEXT DEFAULT NULL
        );
    ''')

    # FTS5 tables (drop and recreate to sync triggers)
    cursor.execute("DROP TABLE IF EXISTS documents_fts")
    cursor.execute("DROP TABLE IF EXISTS chunks_fts")

    cursor.executescript('''
        CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
            title, author, category, tags, full_text,
            content=documents,
            content_rowid=id
        );

        CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
            content,
            content=chunks,
            content_rowid=id
        );

        -- Triggers to keep FTS in sync
        CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN
            INSERT INTO documents_fts(rowid, title, author, category, tags, full_text)
            VALUES (new.id, new.title, new.author, new.category, new.tags, new.full_text);
        END;

        CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN
            INSERT INTO documents_fts(documents_fts, rowid, title, author, category, tags, full_text)
            VALUES ('delete', old.id, old.title, old.author, old.category, old.tags, old.full_text);
        END;

        CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN
            INSERT INTO documents_fts(documents_fts, rowid, title, author, category, tags, full_text)
            VALUES ('delete', old.id, old.title, old.author, old.category, old.tags, old.full_text);
            INSERT INTO documents_fts(rowid, title, author, category, tags, full_text)
            VALUES (new.id, new.title, new.author, new.category, new.tags, new.full_text);
        END;

        CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
            INSERT INTO chunks_fts(rowid, content)
            VALUES (new.id, new.content);
        END;

        CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
            INSERT INTO chunks_fts(chunks_fts, rowid, content)
            VALUES ('delete', old.id, old.content);
        END;

        CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks BEGIN
            INSERT INTO chunks_fts(chunks_fts, rowid, content)
            VALUES ('delete', old.id, old.content);
            INSERT INTO chunks_fts(rowid, content)
            VALUES (new.id, new.content);
        END;
    ''')

    # Rebuild FTS indexes to populate with existing data
    cursor.execute("INSERT INTO documents_fts(documents_fts) VALUES('rebuild')")
    cursor.execute("INSERT INTO chunks_fts(chunks_fts) VALUES('rebuild')")

    conn.commit()
    conn.close()


def extract_text(pdf_path):
    """Extract full text from PDF using PyMuPDF"""
    text_parts = []
    try:
        doc = fitz.open(pdf_path)
        for page in doc:
            text_parts.append(page.get_text())
        doc.close()
    except Exception as e:
        return f"[Error extracting text: {e}]"
    return '\n'.join(text_parts)


def extract_metadata(pdf_path):
    """Extract metadata from PDF"""
    meta = {
        'title': '',
        'author': '',
        'pages': 0,
        'file_size': 0,
        'created_date': ''
    }
    try:
        meta['file_size'] = os.path.getsize(pdf_path)
        doc = fitz.open(pdf_path)
        meta['pages'] = len(doc)
        pdf_meta = doc.metadata
        if pdf_meta:
            meta['title'] = pdf_meta.get('title', '') or ''
            meta['author'] = pdf_meta.get('author', '') or ''
            # Creation date
            cdate = pdf_meta.get('creationDate', '') or ''
            if cdate:
                # Parse PDF date format: D:YYYYMMDDHHmmSS
                match = re.search(r'D:(\d{4})(\d{2})(\d{2})', cdate)
                if match:
                    meta['created_date'] = f"{match.group(1)}-{match.group(2)}-{match.group(3)}"
        doc.close()
    except Exception:
        pass
    return meta


def create_chunks(text, doc_id, chunk_size=800):
    """Split text into semantic chunks of ~chunk_size words"""
    if not text or len(text.strip()) < 50:
        return []

    # Split on paragraph breaks first
    paragraphs = re.split(r'\n\s*\n', text)
    chunks = []
    current_chunk = []
    current_length = 0
    char_start = 0

    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
        words = para.split()
        word_count = len(words)

        if current_length + word_count > chunk_size and current_chunk:
            chunk_text = '\n\n'.join(current_chunk)
            char_end = char_start + len(chunk_text)
            chunks.append({
                'document_id': doc_id,
                'chunk_index': len(chunks),
                'content': chunk_text,
                'char_start': char_start,
                'char_end': char_end
            })
            char_start = char_end
            current_chunk = [para]
            current_length = word_count
        else:
            current_chunk.append(para)
            current_length += word_count

    # Last chunk
    if current_chunk:
        chunk_text = '\n\n'.join(current_chunk)
        char_end = char_start + len(chunk_text)
        chunks.append({
            'document_id': doc_id,
            'chunk_index': len(chunks),
            'content': chunk_text,
            'char_start': char_start,
            'char_end': char_end
        })

    return chunks


def save_to_db(doc, chunks):
    """Save document and chunks to SQLite"""
    conn = get_db()
    cursor = conn.cursor()

    # Check if already indexed
    existing = cursor.execute(
        'SELECT id FROM documents WHERE filepath = ?', (doc['filepath'],)
    ).fetchone()

    if existing:
        # Update existing
        doc_id = existing['id']
        cursor.execute('DELETE FROM chunks WHERE document_id = ?', (doc_id,))
        cursor.execute('''
            UPDATE documents SET
                filename=?, title=?, author=?, pages=?, file_size=?,
                created_date=?, indexed_date=?, category=?, tags=?, full_text=?
            WHERE id=?
        ''', (
            doc['filename'], doc['title'], doc['author'], doc['pages'],
            doc['file_size'], doc['created_date'], doc['indexed_date'],
            doc['category'], doc.get('tags', ''), doc['full_text'], doc_id
        ))
    else:
        cursor.execute('''
            INSERT INTO documents (filename, filepath, title, author, pages, file_size,
                created_date, indexed_date, category, tags, full_text)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            doc['filename'], doc['filepath'], doc['title'], doc['author'],
            doc['pages'], doc['file_size'], doc['created_date'], doc['indexed_date'],
            doc['category'], doc.get('tags', ''), doc['full_text']
        ))
        doc_id = cursor.lastrowid

    # Insert chunks
    for chunk in chunks:
        cursor.execute('''
            INSERT INTO chunks (document_id, chunk_index, content, char_start, char_end)
            VALUES (?, ?, ?, ?, ?)
        ''', (doc_id, chunk['chunk_index'], chunk['content'],
              chunk['char_start'], chunk['char_end']))

    conn.commit()
    conn.close()
    return doc_id


def index_file(pdf_path, original_filename=None):
    """Index a single PDF file. Returns dict with success, category, pages, error."""
    filename = original_filename or os.path.basename(pdf_path)
    try:
        text = extract_text(pdf_path)
        meta = extract_metadata(pdf_path)
        title = meta['title'] or filename.replace('.pdf', '').strip()
        title = re.sub(r'\s*\(AI-\d+\)\s*', '', title).strip()
        category = categorize(filename)

        doc = {
            'filename': filename,
            'filepath': pdf_path,
            'title': title,
            'author': meta['author'],
            'pages': meta['pages'],
            'file_size': meta['file_size'],
            'created_date': meta['created_date'],
            'indexed_date': datetime.now().strftime('%Y-%m-%d %H:%M'),
            'category': category,
            'tags': category.lower(),
            'full_text': text
        }
        chunks = create_chunks(text, 0)
        doc_id = save_to_db(doc, chunks)
        return {'success': True, 'doc_id': doc_id, 'category': category, 'pages': meta['pages'], 'title': title}
    except Exception as e:
        return {'success': False, 'error': str(e)}


def index_directory(path, recursive=True):
    """Index all PDFs in a directory"""
    if not os.path.isdir(path):
        return {'error': f'Directory not found: {path}', 'indexed': 0, 'total': 0}

    pdf_files = []
    if recursive:
        for root, dirs, files in os.walk(path):
            for f in sorted(files):
                if f.lower().endswith('.pdf'):
                    pdf_files.append(os.path.join(root, f))
    else:
        for f in sorted(os.listdir(path)):
            if f.lower().endswith('.pdf'):
                pdf_files.append(os.path.join(path, f))

    indexed = 0
    errors = 0
    results = []

    for i, pdf_path in enumerate(pdf_files, 1):
        filename = os.path.basename(pdf_path)
        try:
            print(f"[{i}/{len(pdf_files)}] Indexing: {filename}")

            # Extract
            text = extract_text(pdf_path)
            meta = extract_metadata(pdf_path)

            # Derive title from filename if metadata empty
            title = meta['title'] or filename.replace('.pdf', '').strip()
            # Clean up title
            title = re.sub(r'\s*\(AI-\d+\)\s*', '', title).strip()

            category = categorize(filename)

            doc = {
                'filename': filename,
                'filepath': pdf_path,
                'title': title,
                'author': meta['author'],
                'pages': meta['pages'],
                'file_size': meta['file_size'],
                'created_date': meta['created_date'],
                'indexed_date': datetime.now().strftime('%Y-%m-%d %H:%M'),
                'category': category,
                'tags': category.lower(),
                'full_text': text
            }

            chunks = create_chunks(text, 0)
            doc_id = save_to_db(doc, chunks)
            indexed += 1
            results.append({'filename': filename, 'status': 'ok', 'doc_id': doc_id})

        except Exception as e:
            errors += 1
            results.append({'filename': filename, 'status': 'error', 'error': str(e)})
            print(f"  ERROR: {e}")

    return {
        'total': len(pdf_files),
        'indexed': indexed,
        'errors': errors,
        'results': results
    }


if __name__ == '__main__':
    init_db()
    result = index_directory('/root/Documents/HSEQ_kennis/')
    print(f"\nDone: {result['indexed']}/{result['total']} indexed, {result['errors']} errors")