"""
Database connection module for Python service
Uses asyncpg for async PostgreSQL connections
"""
import os
import asyncpg
from typing import Optional
from dotenv import load_dotenv

load_dotenv()

# Database connection pool
_pool: Optional[asyncpg.Pool] = None


async def get_pool() -> asyncpg.Pool:
    """Get or create database connection pool"""
    global _pool
    
    if _pool is None:
        # Check for local database first
        use_local_db = os.getenv('USE_LOCAL_DB', 'false').lower() == 'true'
        
        if use_local_db:
            config = {
                'host': os.getenv('LOCAL_DB_HOST', '100.121.163.23'),
                'port': int(os.getenv('LOCAL_DB_PORT', '5432')),
                'user': os.getenv('LOCAL_DB_USER', 'postgres'),
                'password': os.getenv('LOCAL_DB_PASSWORD', 'postgres'),
                'database': os.getenv('LOCAL_DB_NAME', 'institutional_trader')
            }
        elif os.getenv('DATABASE_URL'):
            # Parse connection string
            import urllib.parse
            url = urllib.parse.urlparse(os.getenv('DATABASE_URL'))
            config = {
                'host': url.hostname,
                'port': url.port or 5432,
                'user': url.username,
                'password': url.password,
                'database': url.path.lstrip('/')
            }
        else:
            raise ValueError(
                'No database configuration found. Set either:\n'
                '  - USE_LOCAL_DB=true with LOCAL_DB_* variables\n'
                '  - OR DATABASE_URL'
            )
        
        _pool = await asyncpg.create_pool(
            **config,
            min_size=5,
            max_size=20,
            command_timeout=60
        )
    
    return _pool


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


async def query(sql: str, *args) -> list:
    """Execute a SQL query and return results"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        return await conn.fetch(sql, *args)


async def query_one(sql: str, *args) -> Optional[dict]:
    """Execute a SQL query and return first result"""
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(sql, *args)
        return dict(row) if row else None