# memory_db.py - TiDB Database Operations
import pymysql
from config import settings

class MemoryDB:
    def __init__(self):
        self.connection = pymysql.connect(
            host=settings.TIDB_HOST,
            port=settings.TIDB_PORT,
            user=settings.TIDB_USER,
            password=settings.TIDB_PASSWORD,
            database=settings.TIDB_DATABASE,
            ssl={'ssl': {'ca': ''}}  # TiDB SSL
        )
        self.create_tables()
    
    def create_tables(self):
        with self.connection.cursor() as cursor:
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS chat_history (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    user_input TEXT,
                    ai_response MEDIUMTEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            self.connection.commit()
    
    def save_conversation(self, user_input, ai_response):
        with self.connection.cursor() as cursor:
            cursor.execute(
                "INSERT INTO chat_history (user_input, ai_response) VALUES (%s, %s)",
                (user_input, ai_response)
            )
            self.connection.commit()
    
    def get_recent_conversations(self, limit=10):
        with self.connection.cursor() as cursor:
            cursor.execute(
                "SELECT user_input, ai_response FROM chat_history ORDER BY created_at DESC LIMIT %s",
                (limit,)
            )
            return cursor.fetchall()