from typing import List, Any, Dict, Optional
from urllib.parse import quote_plus
from ..core.base import Sink
from ..common.logger import get_logger
from sqlalchemy import create_engine, text
import pandas as pd

logger = get_logger(__name__)


class PostgreSqlSink(Sink):
    def __init__(self, table: str, 
                 dsn: Optional[str] = None,
                 host: Optional[str] = None, 
                 port: int = 5432, 
                 user: Optional[str] = None, 
                 password: Optional[str] = None, 
                 database: Optional[str] = None,
                 batch_size: int = 1000) -> None:
        self.table = table
        self.batch_size = batch_size
        
        if dsn:
            self.dsn = dsn
        elif host and user and database:
            # Construct DSN with URL-encoded password
            pwd = quote_plus(password) if password else ""
            self.dsn = f"postgresql+psycopg2://{user}:{pwd}@{host}:{port}/{database}"
        else:
            raise ValueError("Must provide either 'dsn' or connection parameters (host, user, password, database)")
            
        self.engine = None
        self._init_engine()

    def _init_engine(self):
        try:
            self.engine = create_engine(self.dsn)
            # Verify connection
            with self.engine.connect() as conn:
                pass
            # Log connection info without exposing credentials
            if "@" in self.dsn:
                host_info = self.dsn.split('@')[-1]
                logger.info(f"Connected to PostgreSQL: {host_info}")
            else:
                logger.info("Connected to PostgreSQL (DSN hidden)")
        except Exception as e:
            logger.error(f"Failed to connect to PostgreSQL: {e}")
            raise

    def write_batch(self, batch: List[Dict[str, Any]]) -> None:
        if not batch:
            return
            
        try:
            df = pd.DataFrame(batch)
            # Use 'append' to add to existing table
            # Use connection context for better transaction handling and compatibility
            with self.engine.begin() as connection:
                df.to_sql(self.table, connection, if_exists='append', index=False, method='multi')
            logger.info(f"Wrote {len(batch)} records to table {self.table}")
        except Exception as e:
            logger.error(f"Failed to write batch to PostgreSQL: {e}")
            raise

    def close(self) -> None:
        if self.engine:
            self.engine.dispose()
