import os
import json
import sqlite3
import requests
import numpy as np
from tqdm import tqdm
import pypdf

# --- CONFIGURATION ---
OLLAMA_URL = "http://localhost:11434/api/embed"
EMBEDDING_MODEL = "qwen3-embedding:0.6b"
DB_PATH = "/root/work/knowledge-base/knowledge_base.db"

KB_JSON_PATH = "/root/work/knowledge-base/juniper-kb/juniper_kb_data_clean.json"
OS_RELEASE_DIR = "/root/work/knowledge-base/juniper-os-release"
BOOKS_DIR = "/root/work/knowledge-base/network-books"

def init_db(db_path):
    """Initializes SQLite database and tables."""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 1. Table for Juniper KB
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS knowledge_base (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        kb_id TEXT NOT NULL,
        salesforce_id TEXT NOT NULL,
        title TEXT,
        url TEXT,
        last_modified TEXT,
        products TEXT,
        categories TEXT,
        environment TEXT,
        symptoms TEXT,
        cause TEXT,
        description TEXT,
        solution TEXT,
        chunk_index INTEGER NOT NULL,
        text_content TEXT NOT NULL,
        prefix_content TEXT NOT NULL,
        embedding BLOB NOT NULL
    )
    """)
    

    
    # 3. Table for Network Books
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS network_book (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        file_name TEXT NOT NULL,
        page_num INTEGER NOT NULL,
        chunk_index INTEGER NOT NULL,
        text_content TEXT NOT NULL,
        prefix_content TEXT NOT NULL,
        embedding BLOB NOT NULL
    )
    """)
    
    # 4. Progress tracker for resumable execution
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS indexing_progress (
        source_type TEXT NOT NULL,
        source_name TEXT NOT NULL,
        status TEXT NOT NULL,
        PRIMARY KEY (source_type, source_name)
    )
    """)
    
    # Create database indices
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_kb_id ON knowledge_base (kb_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_book_file ON network_book (file_name)")
    
    conn.commit()
    return conn

def get_embeddings(texts):
    """Fetches embeddings from the local Ollama API for a batch of texts."""
    try:
        response = requests.post(OLLAMA_URL, json={
            "model": EMBEDDING_MODEL,
            "input": texts
        }, timeout=120)
        response.raise_for_status()
        data = response.json()
        return data.get("embeddings", [])
    except Exception as e:
        print(f"Error calling Ollama API: {e}")
        raise e

def batch_list(lst, batch_size):
    """Helper to split a list into batches."""
    for i in range(0, len(lst), batch_size):
        yield lst[i:i + batch_size]

def chunk_text_words(text, chunk_size_words=300, overlap_words=30):
    """Splits text into chunks of given size using word boundaries."""
    if not text:
        return []
    words = text.split()
    if not words:
        return []
    chunks = []
    i = 0
    while i < len(words):
        chunk_words = words[i:i+chunk_size_words]
        chunks.append(" ".join(chunk_words))
        if i + chunk_size_words >= len(words):
            break
        i += chunk_size_words - overlap_words
    return chunks

def split_markdown_by_headings(text):
    """Splits markdown content into sections based on H2 headings."""
    import re
    parts = re.split(r'^##\s+(.*)$', text, flags=re.MULTILINE)
    sections = []
    
    # Text before the first H2 heading
    intro = parts[0].strip()
    if intro:
        sections.append(("general", intro))
        
    for i in range(1, len(parts), 2):
        heading = parts[i].strip()
        content = parts[i+1].strip() if i+1 < len(parts) else ""
        sections.append((heading, content))
        
    return sections

def process_kb(db_path, json_path, limit=None):
    """Processes and embeds Juniper KB articles from JSON and markdown files."""
    print(f"\n[1/3] Processing Juniper Knowledge Base articles from {json_path}...")
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    if not os.path.exists(json_path):
        print(f"Warning: JSON file not found at {json_path}")
        conn.close()
        return

    with open(json_path, 'r', encoding='utf-8') as f:
        articles = json.load(f)
        
    count = 0
    for article in tqdm(articles, desc="KB Articles"):
        salesforce_id = article.get("kb_id")
        title = article.get("title", "")
        
        # Extract human-readable KB/TN ID from title
        import re
        m = re.search(r'\b((?:KB|TN)\d+)\b', title)
        kb_id = m.group(1) if m else salesforce_id
        
        # Resumability check
        cursor.execute("SELECT status FROM indexing_progress WHERE source_type='kb' AND source_name=?", (salesforce_id,))
        row = cursor.fetchone()
        if row and row[0] == 'completed':
            continue
            
        url = article.get("url", "")
        last_modified = article.get("last_modified", "")
        products = ", ".join(article.get("products", [])) if isinstance(article.get("products"), list) else str(article.get("products") or "")
        categories = ", ".join(article.get("categories", [])) if isinstance(article.get("categories"), list) else str(article.get("categories") or "")
        environment = article.get("environment", "") or ""
        
        # Read raw markdown file if available, to get headers
        file_path = f"/root/work/knowledge-base/juniper-kb/kb_markdown/{salesforce_id}.md"
        markdown_content = ""
        if os.path.exists(file_path):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    markdown_content = f.read()
            except Exception:
                markdown_content = article.get("markdown_content", "")
        else:
            markdown_content = article.get("markdown_content", "")
            
        # Extract semantic sections
        sections = split_markdown_by_headings(markdown_content)
        
        desc_text = ""
        trigger_text = ""
        solution_text = ""
        fixed_text = ""
        symptoms_text = ""
        cause_text = ""
        
        for heading, content in sections:
            h_lower = heading.lower()
            if 'description' in h_lower:
                desc_text = content
            elif 'trigger' in h_lower:
                trigger_text = content
            elif 'solution' in h_lower or 'workaround' in h_lower or 'resolution' in h_lower:
                solution_text = (solution_text + "\n\n" + content) if solution_text else content
            elif 'fixed' in h_lower or 'upgrade' in h_lower or 'fix version' in h_lower or 'fix_version' in h_lower:
                fixed_text = (fixed_text + "\n\n" + content) if fixed_text else content
            elif 'symptom' in h_lower or 'issue' in h_lower:
                symptoms_text = (symptoms_text + "\n\n" + content) if symptoms_text else content
            elif 'cause' in h_lower:
                cause_text = content
                
        # Fallbacks to JSON fields if not extracted from markdown
        if not desc_text:
            desc_text = article.get("description", "") or ""
        if not solution_text:
            solution_text = article.get("solution", "") or ""
        if not symptoms_text:
            symptoms_text = article.get("symptoms", "") or ""
        if not cause_text:
            cause_text = article.get("cause", "") or ""
            
        # PRIORITIZE: Construct description that includes description, trigger, solution, and fixed info
        desc_parts = []
        if desc_text:
            desc_parts.append(desc_text)
        if trigger_text:
            desc_parts.append(f"Trigger: {trigger_text}")
        if solution_text:
            desc_parts.append(f"Solution/Workaround: {solution_text}")
        if fixed_text:
            desc_parts.append(f"Fixed/Upgrade: {fixed_text}")
            
        prioritized_description = "\n\n".join(desc_parts) if desc_parts else desc_text
        
        # Build chunks based on sections
        all_kb_chunks = []
        for heading, content in sections:
            if not content or len(content.strip()) < 5:
                continue
            # Split section content into chunks
            chunks = chunk_text_words(content, chunk_size_words=300, overlap_words=30)
            for idx, chunk in enumerate(chunks):
                prefix = f"[KB Article: {title} (ID: {kb_id}) - {heading}]"
                prefix_chunk = f"{prefix}\n{chunk}"
                all_kb_chunks.append((chunk, prefix_chunk))
                
        if not all_kb_chunks:
            # Fallback if no sections or empty
            fallback_text = f"Title: {title}\nDescription: {prioritized_description}"
            all_kb_chunks = [(fallback_text, f"[KB Article: {title} (ID: {kb_id})]\n{fallback_text}")]
            
        try:
            # Get embeddings in batches
            embeddings = []
            prefix_texts = [item[1] for item in all_kb_chunks]
            
            for batch in batch_list(prefix_texts, 32):
                embeddings.extend(get_embeddings(batch))
                
            # Insert into database
            for idx, ((chunk, prefix_chunk), emb) in enumerate(zip(all_kb_chunks, embeddings)):
                emb_blob = np.array(emb, dtype=np.float32).tobytes()
                cursor.execute("""
                INSERT INTO knowledge_base (
                    kb_id, salesforce_id, title, url, last_modified, products, categories, 
                    environment, symptoms, cause, description, solution, 
                    chunk_index, text_content, prefix_content, embedding
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    kb_id, salesforce_id, title, url, last_modified, products, categories,
                    environment, symptoms_text, cause_text, prioritized_description, solution_text,
                    idx, chunk, prefix_chunk, emb_blob
                ))
                
            cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES ('kb', ?, 'completed')", (salesforce_id,))
            conn.commit()
            
            count += 1
            if limit and count >= limit:
                print(f"Reached KB test limit of {limit}")
                break
                
        except Exception as e:
            print(f"\nError processing KB {kb_id} (SF: {salesforce_id}): {e}")
            conn.rollback()
            cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES ('kb', ?, 'failed')", (salesforce_id,))
            conn.commit()
            
    conn.close()

def process_pdfs(db_path, dir_path, source_type, limit_files=None):
    """Processes PDF files (release notes or books) using optimized batching."""
    print(f"\n[2/3 & 3/3] Processing PDF documents in {dir_path} ({source_type})...")
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    if not os.path.exists(dir_path):
        print(f"Warning: Directory not found at {dir_path}")
        conn.close()
        return

    files = [f for f in os.listdir(dir_path) if f.lower().endswith('.pdf')]
    files.sort()
    
    count_files = 0
    for file in files:
        # Resumability check
        cursor.execute("SELECT status FROM indexing_progress WHERE source_type=? AND source_name=?", (source_type, file))
        row = cursor.fetchone()
        if row and row[0] == 'completed':
            print(f"Skipping already processed PDF: {file}")
            continue
            
        file_path = os.path.join(dir_path, file)
        print(f"\nIndexing PDF file: {file}")
        
        try:
            import pdfplumber
            with pdfplumber.open(file_path) as pdf:
                num_pages = len(pdf.pages)
                
                # Step 1: Extract and chunk all pages
                all_chunks = []
                for page_idx in tqdm(range(num_pages), desc=f"Reading pages of {file}"):
                    page = pdf.pages[page_idx]
                    page_num = page_idx + 1
                    try:
                        text = page.extract_text(x_tolerance=1.5) or ""
                    except Exception:
                        text = page.extract_text() or ""
                    text = text.strip()
                    if not text:
                        continue
                        
                    chunks = chunk_text_words(text, chunk_size_words=300, overlap_words=30)
                    for idx, chunk in enumerate(chunks):
                        prefix = f"[Doc: {file}, Page: {page_num}]"
                        prefix_chunk = f"{prefix}\n{chunk}"
                        all_chunks.append((page_num, idx, chunk, prefix_chunk))
            
            if not all_chunks:
                print(f"No text extracted from PDF: {file}")
                cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES (?, ?, 'completed')", (source_type, file))
                conn.commit()
                continue
                
            # Step 2: Batch embed
            print(f"Generating embeddings for {len(all_chunks)} chunks...")
            embeddings = []
            prefix_texts = [item[3] for item in all_chunks]
            
            for batch in tqdm(batch_list(prefix_texts, 32), total=(len(prefix_texts) + 31) // 32, desc="Embedding batches"):
                embeddings.extend(get_embeddings(batch))
                
            # Step 3: Insert into SQLite in a single transaction
            table_name = "os_release_note" if source_type == "release_note" else "network_book"
            for (page_num, idx, chunk, prefix_chunk), emb in zip(all_chunks, embeddings):
                emb_blob = np.array(emb, dtype=np.float32).tobytes()
                cursor.execute(f"""
                INSERT INTO {table_name} (
                    file_name, page_num, chunk_index, text_content, prefix_content, embedding
                ) VALUES (?, ?, ?, ?, ?, ?)
                """, (
                    file, page_num, idx, chunk, prefix_chunk, emb_blob
                ))
                    
            cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES (?, ?, 'completed')", (source_type, file))
            conn.commit()
            print(f"Successfully indexed PDF: {file}")
            
            count_files += 1
            if limit_files and count_files >= limit_files:
                print(f"Reached {source_type} test file limit of {limit_files}")
                break
                
        except Exception as e:
            print(f"\nError processing PDF {file}: {e}")
            conn.rollback()
            cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES (?, ?, 'failed')", (source_type, file))
            conn.commit()
            
    conn.close()

if __name__ == "__main__":
    import sys
    
    # Clean up previous database in test runs if clean DB is needed
    # (Optional: uncomment if you want a fresh DB for each run, but we keep it for resumability)
    
    # Initialize DB
    init_db(DB_PATH)
    
    # Quick CLI test flag
    is_test = len(sys.argv) > 1 and sys.argv[1] == "--test"
    
    kb_limit = 5 if is_test else None
    pdf_limit = 1 if is_test else None
    
    if is_test:
        print("!!! RUNNING IN TEST MODE (subset of data only) !!!")
        
    try:
        # Process 1: Juniper KB
        process_kb(DB_PATH, KB_JSON_PATH, limit=kb_limit)
        
        # Process 3: Network Books
        process_pdfs(DB_PATH, BOOKS_DIR, "network_book", limit_files=pdf_limit)
        
        print("\nAll done!")
        
    except KeyboardInterrupt:
        print("\nProcess interrupted by user. SQLite transactions rolled back safely.")
