import aiomysql
import asyncio
from secret_handler import return_credentials
import os
from logger_handler import setup_logger


db_username = return_credentials("/etc/secrets/db_username")
db_password = return_credentials("/etc/secrets/db_password")
db_host = os.getenv("BACKEND_PN_DB_HOST","localhost")
db_database = os.getenv("BACKEND_PN_DB_DATABASE","app")

logger = setup_logger(__name__)


class DBManager:
	def __init__(self, host, user, password, db, port=3306, pool_size=5, health_interval=60):
		self.host = host
		self.user = user
		self.password = password
		self.db = db
		self.port = port
		self.pool_size = pool_size
		self._pool: aiomysql.Pool | None = None
		self._health_interval = health_interval
		self._health_task: asyncio.Task | None = None
		self._closing = False

	async def connect(self):
		self._pool = await aiomysql.create_pool(
			host=self.host,
			user=self.user,
			password=self.password,
			db=self.db,
			port=self.port,
			minsize=1,
			maxsize=self.pool_size,
			autocommit=True,
			cursorclass=aiomysql.DictCursor
		)
		logger.info("[DB] Connection pool created")
		self._health_task = asyncio.create_task(self._healthcheck_loop())

	async def _healthcheck_loop(self):
		while not self._closing:
			await asyncio.sleep(self._health_interval)
			try:
				async with self.acquire() as conn:
					async with conn.cursor() as cur:
						await cur.execute("SELECT 1")
				logger.debug("[DB] Healthcheck OK")
			except Exception as e:
				logger.warning(f"[DB] Healthcheck failed: {e}")

	async def acquire(self):
		if not self._pool:
			raise RuntimeError("DB pool not initialized")
		return await self._pool.acquire()

	async def release(self, conn):
		if self._pool:
			self._pool.release(conn)

	async def execute(self, query, *args, retries=3):
		for attempt in range(1, retries + 1):
			conn = await self.acquire()
			try:
				async with conn.cursor() as cur:
					await cur.execute(query, args)
					if cur.description:
						return await cur.fetchall()
					return None
			except aiomysql.OperationalError as e:
				logger.warning(f"[DB] Query failed (attempt {attempt}/{retries}): {e}")
				await asyncio.sleep(2 ** (attempt - 1))
			finally:
				await self.release(conn)
		raise RuntimeError("DB query failed after retries")

	async def close(self):
		self._closing = True
		if self._health_task and not self._health_task.done():
			self._health_task.cancel()
			try:
				await self._health_task
			except asyncio.CancelledError:
				pass
		if self._pool:
			self._pool.close()
			await self._pool.wait_closed()
			logger.info("[DB] Connection pool closed")

db_manager = DBManager(host=db_host, port=3306, user=db_username, password=db_password, db=db_database)