import psycopg2
from psycopg2.extras import RealDictCursor
from .base import BaseConnector
from infer_db.utils.logger_util import logger

class PostgresConnector(BaseConnector):
    def __init__(self, config: dict):
        self.config = config
        self.conn = None
        # Default to 'public' if not specified
        self.target_schema = config.get('schema', 'public')
        
        self._cache = {
            "stats": {},      
            "row_counts": {}, 
            "metadata": {},   
            "existing_constraints": [], 
            "fks": {}        
        }

    def connect(self):
        # We exclude 'schema' from kwargs because psycopg2 doesn't accept it
        conn_args = {k: v for k, v in self.config.items() if k != 'schema'}
        self.conn = psycopg2.connect(**conn_args)

    def get_tables(self) -> list[str]:
        query = """
            SELECT table_name 
            FROM information_schema.tables 
            WHERE table_schema = %s 
            AND table_type = 'BASE TABLE';
        """
        with self.conn.cursor() as cur:
            # FIX: Pass self.target_schema
            cur.execute(query, (self.target_schema,))
            return [row[0] for row in cur.fetchall()]

    def get_columns_metadata(self, table_name: str) -> list[dict]:
        if table_name in self._cache['metadata']:
            return self._cache['metadata'][table_name]
        
        query = """
            SELECT 
                column_name as column, 
                data_type as dtype, 
                is_nullable,
                character_maximum_length as char_len,
                column_default
            FROM information_schema.columns
            WHERE table_name = %s
              AND table_schema = %s  -- FIX: Filter by schema
            ORDER BY ordinal_position;
        """
        with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(query, (table_name, self.target_schema))
            result = [dict(row) for row in cur.fetchall()]
            self._cache['metadata'][table_name] = result
            return result
    
    def get_existing_constraints(self) -> list[dict]:
        if self._cache['existing_constraints']:
            return self._cache['existing_constraints']
        
        query = """
            SELECT tc.table_name, kcu.column_name, tc.constraint_type
            FROM information_schema.table_constraints AS tc
            JOIN information_schema.key_column_usage AS kcu
              ON tc.constraint_name = kcu.constraint_name
              AND tc.table_schema = kcu.table_schema -- Ensure join is within schema
            WHERE tc.constraint_type IN ('PRIMARY KEY', 'FOREIGN KEY')
              AND tc.table_schema = %s; -- FIX: Filter by schema
        """
        with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(query, (self.target_schema,))
            result = [dict(row) for row in cur.fetchall()]
            self._cache['existing_constraints'] = result
            return result

    def get_column_stats(self, table_name: str, column_name: str) -> dict:
        """
        Returns a dict: {'min': val, 'max': val, 'distinct': int, 'nulls': int}
        """
        cache_key = f"{table_name}.{column_name}"
        if cache_key in self._cache['stats']:
            return self._cache['stats'][cache_key]

        # LOGIC:
        # COUNT(*) = Total Rows
        # COUNT("col") = Non-Null Rows
        # Difference = Null Rows
        query = f"""
            SELECT 
                MIN("{column_name}"), 
                MAX("{column_name}"), 
                COUNT(DISTINCT "{column_name}"),
                (COUNT(*) - COUNT("{column_name}")) AS null_count
            FROM "{self.target_schema}"."{table_name}"
        """
        try:
            with self.conn.cursor() as cur:
                cur.execute(query)
                row = cur.fetchone()
                
                # Convert to a dictionary for safer access
                stats = {
                    "min": row[0],
                    "max": row[1],
                    "distinct": row[2],
                    "nulls": row[3] # <--- NEW DATA POINT
                }
                
                self._cache['stats'][cache_key] = stats
                return stats
        except Exception as e:
            logger.error(f"Could not fetch stats for {table_name}.{column_name}: {e}")
            self.conn.rollback()
            return None
    
    def get_row_count(self, table_name: str) -> int:
        if table_name in self._cache['row_counts']:
            return self._cache['row_counts'][table_name]
        
        # FIX: Fully qualify table name
        query = f'SELECT COUNT(*) FROM "{self.target_schema}"."{table_name}";'
        try:
            with self.conn.cursor() as cur:
                cur.execute(query)
                result = cur.fetchone()
                count = result[0] if result else 0
                self._cache['row_counts'][table_name] = count
                return count
        except Exception as e:
            logger.error(f"Error fetching row count for {table_name}: {e}")
            self.conn.rollback()
            return 0
    
    def get_existing_fks(self, table_name):
        if table_name in self._cache["fks"]:
            return self._cache["fks"][table_name]

        query = """
            SELECT
                kcu.column_name AS source_col,
                ccu.table_name AS target_table,
                ccu.column_name AS target_col
            FROM information_schema.key_column_usage AS kcu
            JOIN information_schema.constraint_column_usage AS ccu
                ON ccu.constraint_name = kcu.constraint_name
                AND ccu.table_schema = kcu.table_schema
            JOIN information_schema.table_constraints AS tc
                ON tc.constraint_name = kcu.constraint_name
                AND tc.table_schema = kcu.table_schema
            WHERE tc.constraint_type = 'FOREIGN KEY'
              AND kcu.table_name = %s
              AND kcu.table_schema = %s; -- FIX: Filter by schema
        """
        try:
            with self.conn.cursor() as cur:
                cur.execute(query, (table_name, self.target_schema))
                results = cur.fetchall()
                
                fks = []
                for row in results:
                    fks.append({
                        "source_col": row[0],
                        "target_table": row[1],
                        "target_col": row[2]
                    })
                
                self._cache["fks"][table_name] = fks
                return fks
                
        except Exception as e:
            logger.error(f"Error fetching FKs for {table_name}: {e}")
            return []
        
    def check_containment(self, source_table: str, source_col: str, target_table: str, target_col: str) -> bool:
        # FIX: Fully qualify both source and target tables
        query = f"""
            SELECT 1 
            FROM "{self.target_schema}"."{source_table}" s
            WHERE s."{source_col}" IS NOT NULL 
              AND NOT EXISTS (
                  SELECT 1 
                  FROM "{self.target_schema}"."{target_table}" t 
                  WHERE t."{target_col}" = s."{source_col}"
              )
            LIMIT 1;
        """
        try:
            with self.conn.cursor() as cur:
                cur.execute(query)
                orphan = cur.fetchone()
                return orphan is None
        except Exception as e:
            self.conn.rollback()
            return False