import database


class TgDatabase(database.BasicDatabase):
    def __init__(self, hostname: str, user: str, password: str, dbname: str):
        super().__init__(hostname, user, password, dbname)

        with self.pool.acquire() as conn:
            with conn.cursor() as cursor:
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS bots (
                    id BIGINT NOT NULL,
                    owner_id BIGINT NOT NULL,
                    api_token VARCHAR(64) NOT NULL,
                    ai_prompt VARCHAR(2000) DEFAULT NULL,
                    group_chats_allowed TINYINT NOT NULL DEFAULT 1,
                    private_chats_allowed TINYINT NOT NULL DEFAULT 1,
                    PRIMARY KEY (id)
                )""")

                cursor.execute("""
                CREATE TABLE IF NOT EXISTS chats (
                    bot_id BIGINT NOT NULL,
                    chat_id BIGINT NOT NULL,
                    active TINYINT NOT NULL DEFAULT 0,
                    rules VARCHAR(4000),
                    greeting_join VARCHAR(2000),
                    ai_prompt VARCHAR(2000),
                    PRIMARY KEY (bot_id, chat_id),
                    CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id)
                        ON UPDATE CASCADE ON DELETE CASCADE
                )""")

                cursor.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    bot_id BIGINT NOT NULL,
                    chat_id	BIGINT NOT NULL,
                    user_id	BIGINT NOT NULL,
                    last_message BIGINT NOT NULL DEFAULT 0,
                    messages_today SMALLINT NOT NULL DEFAULT 0,
                    messages_month SMALLINT NOT NULL DEFAULT 0,
                    warnings TINYINT NOT NULL DEFAULT 0,
                    about VARCHAR(1000),
                    PRIMARY KEY (bot_id, chat_id, user_id),
                    CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id)
                        ON UPDATE CASCADE ON DELETE CASCADE
                )""")

                cursor.execute("""
                CREATE TABLE IF NOT EXISTS contexts (
                    id BIGINT NOT NULL auto_increment,
                    bot_id BIGINT NOT NULL,
                    chat_id BIGINT NOT NULL,
                    message_id BIGINT,
                    role VARCHAR(16) NOT NULL,
                    text VARCHAR(4000),
                    image MEDIUMBLOB,
                    PRIMARY KEY (id),
                    CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id)
                        ON UPDATE CASCADE ON DELETE CASCADE
                )""")


DB: TgDatabase


def create_database(hostname: str, user: str, password: str, dbname: str):
    global DB
    DB = TgDatabase(hostname, user, password, dbname)