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

class PostgresConnector(BaseConnector):
    def __init__(self, config: dict):
        self.config = config # dict with host, dbname, user, password
        self.conn = None
        self._cache = {
            "stats": {},      # { 'table.column': (min, max, distinct) }
            "row_counts": {}, # { 'table': count }
            "metadata": {},    # { 'table': [columns...] }
            "existing_constraints": [], # Cached existing constraints
            "fks": {}        # { 'table_name': [ {fk_info}... ] }
        }

    def connect(self):
        self.conn = psycopg2.connect(**self.config)

    def get_tables(self) -> list[str]:
        query = """
            SELECT table_name 
            FROM information_schema.tables 
            WHERE table_schema = 'public';
        """
        with self.conn.cursor() as cur:
            cur.execute(query)
            return [row[0] for row in cur.fetchall()]

    def get_columns_metadata(self, table_name: str) -> list[dict]:
        # Caches metadata per table
        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
            ORDER BY ordinal_position;
        """
        with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(query, (table_name,))
            # Convert RealDictRow to standard dict for better compatibility
            result = [dict(row) for row in cur.fetchall()]
            self._cache['metadata'][table_name] = result
            return result
    
    def get_existing_constraints(self) -> list[dict]:
        # Queries existing PK/FK constraints so we don't "re-detect" them
        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
            WHERE tc.constraint_type IN ('PRIMARY KEY', 'FOREIGN KEY');
        """
        with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(query)
            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) -> tuple[int, int, int]:
        """Returns (min, max, count_distinct)."""
        cache_key = f"{table_name}.{column_name}"
        if cache_key in self._cache['stats']:
            return self._cache['stats'][cache_key]
        # We cast to numeric for min/max to handle various ID types (int, bigint)
        query = f"""
            SELECT 
                MIN("{column_name}"), 
                MAX("{column_name}"), 
                COUNT(DISTINCT "{column_name}") 
            FROM "{table_name}"
        """
        try:
            with self.conn.cursor() as cur:
                cur.execute(query)
                result = cur.fetchone()
                self._cache['stats'][cache_key] = result
                return result
        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:
        """Returns the total number of rows in a table."""
        if table_name in self._cache['row_counts']:
            return self._cache['row_counts'][table_name]
        
        query = f'SELECT COUNT(*) FROM "{table_name}";'
        try:
            with self.conn.cursor() as cur:
                cur.execute(query)
                result = cur.fetchone()
                self._cache['row_counts'][table_name] = result[0] if result else 0
                return self._cache['row_counts'][table_name]
        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):
        """
        Returns a list of explicit foreign keys for the given table.
        Uses cache to avoid redundant queries.
        """
        # 1. Check Cache
        if table_name in self._cache["fks"]:
            return self._cache["fks"][table_name]

        # 2. Fetch if missing
        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
            JOIN information_schema.table_constraints AS tc
                ON tc.constraint_name = kcu.constraint_name
            WHERE tc.constraint_type = 'FOREIGN KEY'
              AND kcu.table_name = %s;
        """
        try:
            with self.conn.cursor() as cur:
                cur.execute(query, (table_name,))
                results = cur.fetchall()
                
                fks = []
                for row in results:
                    fks.append({
                        "source_col": row[0],
                        "target_table": row[1],
                        "target_col": row[2]
                    })
                
                # 3. Store in Cache
                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:
        """
        Returns True if ALL non-null values in source_col exist in target_col.
        Uses a LEFT JOIN / NOT EXISTS query for speed.
        """
        # Query: Find ONE row in Source that has NO match in Target
        query = f"""
            SELECT 1 
            FROM "{source_table}" s
            WHERE s."{source_col}" IS NOT NULL 
              AND NOT EXISTS (
                  SELECT 1 
                  FROM "{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()
                # If we found an orphan (result exists), Containment is FALSE.
                # If orphan is None, Containment is TRUE.
                return orphan is None
        except Exception as e:
            # If types are strictly incompatible (e.g. comparing int to string without cast),
            # Postgres might throw an error. We treat this as "Not a Match".
            # CRITICAL FIX: If a query fails (e.g. really bad types), the transaction dies.
            # We MUST rollback to reset the connection for the next test.
            self.conn.rollback()
            return False