database.py

import aiomysql
import os
from datetime import datetime
from typing import Optional

# MySQL connection pool
_pool: Optional[aiomysql.Pool] = None


async def get_pool() -> aiomysql.Pool:
    """Get or create MySQL connection pool"""
    global _pool
    if _pool is None:
        _pool = await aiomysql.create_pool(
            host=os.getenv("MYSQL_HOST", "localhost"),
            port=int(os.getenv("MYSQL_PORT", 3306)),
            user=os.getenv("MYSQL_USER", "root"),
            password=os.getenv("MYSQL_PASSWORD", ""),
            db=os.getenv("MYSQL_DATABASE", "vtorynka_bot"),
            autocommit=True,
            charset="utf8mb4",
            minsize=1,
            maxsize=10
        )
    return _pool


async def close_pool():
    """Close the connection pool"""
    global _pool
    if _pool:
        _pool.close()
        await _pool.wait_closed()
        _pool = None


async def init_db():
    """Initialize database"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("""
                CREATE TABLE IF NOT EXISTS conversations (
                    user_id BIGINT PRIMARY KEY,
                    username VARCHAR(255),
                    language VARCHAR(10) DEFAULT NULL,
                    history LONGTEXT,
                    reminder_count INT DEFAULT 0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
            """)

            # Migration: add language column if it doesn't exist
            try:
                await cur.execute("""
                    ALTER TABLE conversations
                    ADD COLUMN language VARCHAR(10) DEFAULT NULL
                """)
            except Exception:
                pass  # Column already exists

            # Migration: add reminder_count column if it doesn't exist
            try:
                await cur.execute("""
                    ALTER TABLE conversations
                    ADD COLUMN reminder_count INT DEFAULT 0
                """)
            except Exception:
                pass  # Column already exists

            await cur.execute("""
                CREATE TABLE IF NOT EXISTS seller_leads (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    user_id BIGINT,
                    username VARCHAR(255),
                    client_name VARCHAR(255),
                    phone VARCHAR(50),
                    property_type VARCHAR(100),
                    address VARCHAR(500),
                    rooms VARCHAR(50),
                    area VARCHAR(50),
                    price VARCHAR(100),
                    callback_time VARCHAR(100),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    status VARCHAR(50) DEFAULT 'new',
                    notes TEXT,
                    INDEX idx_user_id (user_id),
                    INDEX idx_status (status),
                    INDEX idx_created (created_at)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
            """)


async def get_user_language(user_id: int) -> Optional[str]:
    """Get user's selected language"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT language FROM conversations WHERE user_id = %s", (user_id,)
            )
            row = await cur.fetchone()
            if row:
                return row[0]
            return None


async def set_user_language(user_id: int, username: Optional[str], language: str):
    """Set user's language preference"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("""
                INSERT INTO conversations (user_id, username, language, updated_at)
                VALUES (%s, %s, %s, %s)
                ON DUPLICATE KEY UPDATE
                    language = VALUES(language),
                    updated_at = VALUES(updated_at)
            """, (user_id, username, language, datetime.now()))


async def get_conversation_history(user_id: int) -> list:
    """Get user's conversation history"""
    import json
    pool = await get_pool()
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute(
                "SELECT history FROM conversations WHERE user_id = %s", (user_id,)
            )
            row = await cur.fetchone()
            if row and row[0]:
                return json.loads(row[0])
            return []


async def save_conversation_history(user_id: int, username: Optional[str], history: list):
    """Save conversation history"""
    import json
    history_json = json.dumps(history, ensure_ascii=False)
    pool = await get_pool()
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("""
                INSERT INTO conversations (user_id, username, history, updated_at)
                VALUES (%s, %s, %s, %s)
                ON DUPLICATE KEY UPDATE
                    history = VALUES(history),
                    updated_at = VALUES(updated_at)
            """, (user_id, username, history_json, datetime.now()))


async def save_seller_lead(
    user_id: int,
    username: Optional[str],
    client_name: str,
    phone: str,
    property_type: str = "",
    address: str = "",
    rooms: str = "",
    area: str = "",
    price: str = "",
    callback_time: str = ""
):
    """Save seller lead for exclusive agreement"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("""
                INSERT INTO seller_leads (
                    user_id, username, client_name, phone,
                    property_type, address, rooms, area, price, callback_time
                )
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """, (
                user_id, username, client_name, phone,
                property_type, address, rooms, area, price, callback_time
            ))