import sqlite3
import json
from langchain.memory import ConversationBufferMemory

DB_FILE = ".conversation_memory.db"

class Memory:
    def __init__(self, memory_key="chat_history"):
        self.memory_key = memory_key
        self.memory = ConversationBufferMemory(memory_key=memory_key, return_messages=True)
        self._init_db()
        self._load_history()

    def _init_db(self):
        conn = sqlite3.connect(DB_FILE)
        cursor = conn.cursor()
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS conversations (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            question TEXT,
            answer TEXT
        )
        """)
        conn.commit()
        conn.close()

    def _load_history(self):
        conn = sqlite3.connect(DB_FILE)
        cursor = conn.cursor()
        cursor.execute("SELECT question, answer FROM conversations")
        rows = cursor.fetchall()
        conn.close()

        for question, answer in rows:
            self.memory.chat_memory.add_user_message(question)
            self.memory.chat_memory.add_ai_message(answer)

    def save_context(self, inputs, outputs):
        """Save both to memory and to SQLite"""
        question = inputs.get("question")
        answer = outputs.get("output")
        if question and answer:
            # Save to in-memory buffer
            self.memory.save_context(inputs, outputs)
            
            # Save to SQLite
            conn = sqlite3.connect(DB_FILE)
            cursor = conn.cursor()
            cursor.execute(
                "INSERT INTO conversations (question, answer) VALUES (?, ?)",
                (question, answer)
            )
            conn.commit()
            conn.close()

    def load_memory_variables(self, inputs=None):
        return self.memory.load_memory_variables(inputs)