import sqlite3
import pandas as pd
from datetime import datetime
import os

DB_NAME = "navy_docs.db"

class DatabaseManager:
    def __init__(self, db_path=DB_NAME):
        self.db_path = db_path
        self._init_db()

    def _init_db(self):
            """Creates tables if they don't exist."""
            conn = sqlite3.connect(self.db_path)
            c = conn.cursor()
            
            # UPDATED: Added 'abstract' column
            c.execute('''
                CREATE TABLE IF NOT EXISTS documents (
                    doc_id TEXT PRIMARY KEY,
                    filename TEXT,
                    upload_date TEXT,
                    user_id TEXT, 
                    full_text TEXT,
                    abstract TEXT 
                )
            ''')
            
            # (Chunks table remains the same...)
            c.execute('''
                CREATE TABLE IF NOT EXISTS chunks (
                    chunk_id TEXT PRIMARY KEY,
                    doc_id TEXT,
                    text TEXT,
                    chunk_index INTEGER,
                    FOREIGN KEY(doc_id) REFERENCES documents(doc_id)
                )
            ''')
            conn.commit()
            conn.close()

    # UPDATED: Added abstract argument
    def add_document(self, doc_id, filename, full_text, abstract="", user_id="default"):
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute(
            "INSERT OR REPLACE INTO documents (doc_id, filename, upload_date, user_id, full_text, abstract) VALUES (?, ?, ?, ?, ?, ?)",
            (doc_id, filename, datetime.now().isoformat(), user_id, full_text, abstract)
        )
        conn.commit()
        conn.close()
        
    # NEW: Getter for the abstract
    def get_doc_abstract(self, doc_id):
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("SELECT abstract FROM documents WHERE doc_id=?", (doc_id,))
        result = c.fetchone()
        conn.close()
        return result[0] if result else "No abstract available."

    def get_doc_text(self, doc_id):
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("SELECT full_text FROM documents WHERE doc_id=?", (doc_id,))
        result = c.fetchone()
        conn.close()
        return result[0] if result else "Error: Document text not found."

    def get_all_filenames(self):
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("SELECT filename FROM documents")
        results = [r[0] for r in c.fetchall()]
        conn.close()
        return results

    def delete_document(self, filename):
        """Deletes a document and its chunks by filename."""
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        
        # 1. Get doc_id
        c.execute("SELECT doc_id FROM documents WHERE filename=?", (filename,))
        res = c.fetchone()
        if not res:
            conn.close()
            return None
            
        doc_id = res[0]
        
        # 2. Delete chunks first (Foreign Key hygiene)
        c.execute("DELETE FROM chunks WHERE doc_id=?", (doc_id,))
        
        # 3. Delete parent
        c.execute("DELETE FROM documents WHERE doc_id=?", (doc_id,))
        
        conn.commit()
        conn.close()
        return doc_id