import sqlite3
import psycopg2
import os
from dotenv import load_dotenv

load_dotenv()

class Database:
    def __init__(self, db_type="sqlite"):
        self.db_type = db_type
        self.connection = None

    def connect(self):
        """Connects to the chosen database (PostgreSQL or SQLite)."""
        if self.db_type == "postgres":
            self.connection = psycopg2.connect(
                dbname=os.getenv("POSTGRES_DB"),
                user=os.getenv("POSTGRES_USER"),
                password=os.getenv("POSTGRES_PASSWORD"),
                host=os.getenv("POSTGRES_HOST"),
                port=os.getenv("POSTGRES_PORT", "5432"),
            )
        else:
            self.connection = sqlite3.connect(os.getenv("SQLITE_DB_PATH", "prompts.db"))
        return self.connection

    def close(self):
        """Closes the database connection."""
        if self.connection:
            self.connection.close()

    def create_table(self):
        """Creates the prompt_configurations table if it does not exist."""
        conn = self.connect()
        cursor = conn.cursor()
        
        create_table_query = """
        CREATE TABLE IF NOT EXISTS prompt_configurations (
            id BIGINT GENERATED BY DEFAULT AS IDENTITY (INCREMENT BY 1 MINVALUE 1 MAXVALUE 9223372036854775807 START 1 CACHE 1 NO CYCLE) PRIMARY KEY,
            prompt_type VARCHAR(255) NOT NULL,
            description TEXT NOT NULL,
            system_message TEXT NOT NULL,
            human_message TEXT NOT NULL,
            structured_output BOOLEAN DEFAULT FALSE,
            output_format TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        """ if self.db_type == "postgres" else """
        CREATE TABLE IF NOT EXISTS prompt_configurations (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            prompt_type TEXT NOT NULL,
            description TEXT NOT NULL,
            system_message TEXT NOT NULL,
            human_message TEXT NOT NULL,
            structured_output BOOLEAN DEFAULT 0,
            output_format TEXT,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        );
        """
        
        cursor.execute(create_table_query)
        conn.commit()
        self.close()
