#!/usr/bin/env python3
"""
SMTP Connection Pool with robust connection handling
"""

import smtplib
from queue import Queue, Empty
from typing import Optional

from logger import log
from config import config


class SMTPPool:
    """Thread-safe SMTP Connection Pool"""
    
    def __init__(self, host: str, port: int, pool_size: int = 5):
        self.host = host
        self.port = port
        self.pool_size = pool_size
        self._pool: Queue = Queue(maxsize=pool_size)
        self._initialized = False
    
    def _create_connection(self) -> Optional[smtplib.SMTP]:
        """Create new SMTP connection"""
        try:
            conn = smtplib.SMTP(self.host, self.port, timeout=30)
            conn.ehlo()
            if config.smtp_use_tls:
                conn.starttls()
                conn.ehlo()
            if config.smtp_user and config.smtp_pass:
                conn.login(config.smtp_user, config.smtp_pass)
            log(f"  📡 New SMTP connection created to {self.host}:{self.port}")
            return conn
        except Exception as e:
            log(f"Failed to create SMTP connection: {e}", 'ERROR')
            return None
    
    def _test_connection(self, conn: smtplib.SMTP) -> bool:
        """Test if connection is still alive"""
        try:
            status = conn.noop()[0]
            return status == 250
        except Exception:
            return False
    
    def initialize(self):
        """Pre-create connections"""
        if self._initialized:
            return
        
        # Only 1-2 connections initially, rest on-demand
        for _ in range(min(2, self.pool_size)):
            conn = self._create_connection()
            if conn:
                self._pool.put(conn)
        
        self._initialized = True
        log(f"SMTP pool initialized with {self._pool.qsize()} connections (max: {self.pool_size})")
    
    def get_connection(self, timeout: float = 5.0) -> Optional[smtplib.SMTP]:
        """Get a valid connection from pool or create new one"""
        # Try to get from pool
        try:
            conn = self._pool.get(block=False)
            # Test if connection is still alive
            if self._test_connection(conn):
                return conn
            else:
                # Connection is dead, close and create new one
                try:
                    conn.quit()
                except:
                    pass
                log(f"  ♻ Recycled stale SMTP connection")
                return self._create_connection()
        except Empty:
            # Pool empty, create new connection
            return self._create_connection()
    
    def return_connection(self, conn: smtplib.SMTP):
        """Return connection to pool if still valid"""
        if conn is None:
            return
        
        # Check if connection is still good
        if not self._test_connection(conn):
            try:
                conn.quit()
            except:
                pass
            log(f"  🗑 Discarded broken SMTP connection")
            return
        
        # Try to return to pool
        try:
            self._pool.put_nowait(conn)
        except:
            # Pool full, close connection
            try:
                conn.quit()
            except:
                pass
    
    def close_all(self):
        """Close all connections"""
        while not self._pool.empty():
            try:
                conn = self._pool.get_nowait()
                conn.quit()
            except:
                pass