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
from sqlalchemy.dialects.postgresql import insert
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,
                 upsert_keys: Optional[List[str]] = None,
                 batch_size: int = 1000) -> None:
        self.table = table
        self.upsert_keys = upsert_keys
        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 _create_upsert_method(self, upsert_keys: List[str]):
        def method(table, conn, keys, data_iter):
            # Convert data iterator to list of dicts
            data = [dict(zip(keys, row)) for row in data_iter]
            
            if not data:
                return

            # Create insert statement
            stmt = insert(table.table).values(data)
            
            # Construct set_ dict: update all columns that are being inserted, except the conflict keys
            set_dict = {
                key: getattr(stmt.excluded, key)
                for key in keys
                if key not in upsert_keys
            }
            
            # Create upsert statement
            upsert_stmt = stmt.on_conflict_do_update(
                index_elements=upsert_keys,
                set_=set_dict
            )
            
            # Execute statement
            conn.execute(upsert_stmt)
            
        return method

    def write_batch(self, batch: List[Dict[str, Any]]) -> None:
        if not batch:
            return
            
        try:
            # Debug log to inspect data
            if logger.isEnabledFor(10): # DEBUG level
                logger.debug(f"Writing batch of {len(batch)} records. First record: {batch[0]}")
                
            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:
                if self.upsert_keys:
                    method = self._create_upsert_method(self.upsert_keys)
                    df.to_sql(self.table, connection, if_exists='append', index=False, method=method)
                else:
                    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()
