import sqlite3
import numpy as np
import pickle
import os

DB_NAME = os.getenv("DB_NAME", "documents.db")


def get_connection():
    conn = sqlite3.connect(DB_NAME, timeout=30)
    conn.execute("PRAGMA journal_mode=WAL;")
    conn.execute("PRAGMA busy_timeout=30000;")
    return conn


def create_tables():
    conn = get_connection()
    cursor = conn.cursor()

    cursor.execute("""
        CREATE TABLE IF NOT EXISTS document_chunks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT,
            page INTEGER,
            chunk_id INTEGER,
            content TEXT,
            embedding BLOB
        )
    """)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS query_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT,
            question TEXT,
            answer TEXT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
        )
        """)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE,
            password TEXT,
            role TEXT
        )
        """)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS documents (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT UNIQUE,
            uploaded_by TEXT,
            uploaded_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
        """)


    conn.commit()
    conn.close()


def insert_chunk(chunk, embedding):
    conn = get_connection()
    cursor = conn.cursor()

    embedding_blob = pickle.dumps(embedding)

    cursor.execute("""
        INSERT INTO document_chunks
        (filename, page, chunk_id, content, embedding)
        VALUES (?, ?, ?, ?, ?)
    """, (
        chunk["filename"],
        chunk["page"],
        chunk["chunk_id"],
        chunk["content"],
        embedding_blob
    ))

    conn.commit()
    conn.close()


def insert_chunks_bulk(chunks, embeddings):
    conn = get_connection()
    cursor = conn.cursor()

    rows = []
    for chunk, embedding in zip(chunks, embeddings):
        rows.append(
            (
                chunk["filename"],
                chunk["page"],
                chunk["chunk_id"],
                chunk["content"],
                pickle.dumps(embedding),
            )
        )

    cursor.executemany(
        """
        INSERT INTO document_chunks
        (filename, page, chunk_id, content, embedding)
        VALUES (?, ?, ?, ?, ?)
        """,
        rows,
    )

    conn.commit()
    conn.close()


def fetch_all_chunks():
    conn = get_connection()
    cursor = conn.cursor()

    cursor.execute("""
        SELECT filename, page, chunk_id, content, embedding
        FROM document_chunks
    """)

    rows = cursor.fetchall()
    conn.close()

    chunks = []
    embeddings = []

    for row in rows:
        chunks.append({
            "filename": row[0],
            "page": row[1],
            "chunk_id": row[2],
            "content": row[3]
        })
        embeddings.append(pickle.loads(row[4]))

    return chunks, np.array(embeddings)

def document_exists(filename):
    conn = get_connection()
    cursor = conn.cursor()

    cursor.execute(
        "SELECT COUNT(*) FROM documents WHERE filename = ?",
        (filename,)
    )

    exists = cursor.fetchone()[0] > 0
    conn.close()

    return exists

def insert_query(username, question, answer):

    conn = get_connection()
    cursor = conn.cursor()

    cursor.execute(
        "INSERT INTO query_history (username, question, answer) VALUES (?, ?, ?)",
        (username, question, answer[:500])
    )

    conn.commit()
    conn.close()

def fetch_query_history(limit=10):
    conn = get_connection()
    cursor = conn.cursor()

    cursor.execute("""
        SELECT question, answer, timestamp
        FROM query_history
        ORDER BY timestamp DESC
        LIMIT ?
    """, (limit,))

    rows = cursor.fetchall()
    conn.close()
    return rows

def clear_query_history():
    conn = get_connection()
    cursor = conn.cursor()

    cursor.execute("DELETE FROM query_history")

    conn.commit()
    conn.close()