from contextlib import asynccontextmanager
import aiomysql
import asyncio
from secret_handler import return_credentials
import os
from logger_handler import setup_logger
import pymysql.err

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
		logger.debug(f"[DB] Initialized DBManager: host={host}, db={db}, port={port}, pool_size={pool_size}")

	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")
		logger.debug(f"[DB] Pool object: {self._pool}")
		self._health_task = asyncio.create_task(self._healthcheck_loop())
		logger.debug("[DB] Healthcheck task started")

	async def _healthcheck_loop(self):
		logger.debug("[DB] Healthcheck loop running")
		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] Pool healthcheck succeeded")
			except pymysql.err.Error as e:
				logger.warning(f"[DB] Healthcheck failed: {e}")
				
	@asynccontextmanager
	async def acquire(self):
		logger.debug("[DB] Acquiring connection from pool")
		conn = await self._pool.acquire()
		logger.debug(f"[DB] Connection acquired: {conn}")
		try:
			yield conn
		finally:
			self._pool.release(conn)
			logger.debug(f"[DB] Connection released: {conn}")

	async def release(self, conn):
		if self._pool:
			self._pool.release(conn)
			logger.debug(f"[DB] Connection manually released: {conn}")

	async def close(self):
		logger.info("[DB] Closing connection pool...")
		self._closing = True
		if self._health_task and not self._health_task.done():
			self._health_task.cancel()
			try:
				await self._health_task
				logger.debug("[DB] Healthcheck task cancelled")
			except asyncio.CancelledError:
				logger.debug("[DB] Healthcheck task cancellation caught")
		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)