# tenant_manager/async/pg_client.py
 
import asyncpg
import logging
from config import DBConfig
from tenant_manage.tenant_exceptions import TenantNotFound, UserNotFound
from tenant_manage.tenant_cache import TenantCache
 
logger = logging.getLogger(__name__)
 
 
class AsyncPostgresClient:
    """
    Handles multi-tenant DB:
    1. Connect to Fyndo main DB
    2. Get tenant DB name from 'organizations' table
    3. Connect to tenant DB
    4. Validate user inside tenant DB
    """
 
    def __init__(self, tenant_id: str):
        self.tenant_id = tenant_id
        self.fyndo_pool = None
        self.tenant_pool = None
 
    # -----------------------------------------
    # 1️⃣ Connect to MAIN FYNDODB (organizations)
    # -----------------------------------------
    async def connect_fyndo(self):
        if self.fyndo_pool:
            return self.fyndo_pool
 
        dsn = DBConfig.build_main_db_dsn()
 
        self.fyndo_pool = await asyncpg.create_pool(
            dsn=dsn,
            min_size=1,
            max_size=5
        )
        return self.fyndo_pool
 
 
 
    # -----------------------------------------
    # 2️⃣ Get tenant database name from org table
    # -----------------------------------------
    async def get_tenant_db_name(self) -> str:
        cached = TenantCache.get(self.tenant_id)
        if cached:
            return cached
 
        pool = await self.connect_fyndo()
 
        query = """
            SELECT db_name
            FROM organizations
            WHERE id = $1
            LIMIT 1
        """
 
 
        async with pool.acquire() as conn:
            row = await conn.fetchrow(query, self.tenant_id)
 
            if not row:
                raise TenantNotFound(f"Tenant {self.tenant_id} not found")
 
            TenantCache.set(self.tenant_id, row["db_name"])
            return row["db_name"]
 
    # -----------------------------------------
    # 3️⃣ Connect to the tenant DB
    # -----------------------------------------
    async def connect_tenant(self):
        if self.tenant_pool:
            return self.tenant_pool
 
        tenant_db_name = await self.get_tenant_db_name()
 
        dsn = DBConfig.build_tenant_db_dsn(tenant_db_name)
 
        # print("LLLLLLLLLLLLLLLLLLLLLLLLLLLLLL",dsn)
 
        self.tenant_pool = await asyncpg.create_pool(
            dsn=dsn,
            min_size=1,
            max_size=5
        )
        return self.tenant_pool
 
    # -----------------------------------------
    # 4️⃣ Validate user inside tenant DB
    # -----------------------------------------
    async def fetch_user(self, user_id: str):
        pool = await self.connect_tenant()
 
        query = """
            SELECT id, role_id
            FROM users
            WHERE id = $1
            LIMIT 1
        """
 
        async with pool.acquire() as conn:
            row = await conn.fetchrow(query, user_id)
 
            if not row:
                raise UserNotFound(f"User {user_id} not found in tenant DB")
 
            return dict(row)
 
   
   
 
        # -----------------------------------------
    # 5️⃣ Get tenant default LLM provider
    # -----------------------------------------
    async def fetch_tenant_default_llm(self) -> str:
        pool = await self.connect_fyndo()
 
        query = """
            SELECT default_llm_provider
            FROM organizations
            WHERE id = $1
            LIMIT 1
        """
 
        async with pool.acquire() as conn:
            row = await conn.fetchrow(query, self.tenant_id)
 
            if not row:
                raise TenantNotFound(f"Tenant {self.tenant_id} not found")
 
            return row["default_llm_provider"]
 
 
 
    async def execute(self, query: str, *args):
        pool = await self.connect_tenant()
        async with pool.acquire() as conn:
            return await conn.execute(query, *args)
 
 
    async def fetchrow(self, query: str, *args):
        pool = await self.connect_tenant()
        async with pool.acquire() as conn:
            return await conn.fetchrow(query, *args)
 
 
    async def fetch(self, query: str, *args):
        pool = await self.connect_tenant()
        async with pool.acquire() as conn:
            return await conn.fetch(query, *args)