from .base import PKStrategy
from infer_db.utils.logger_util import logger
from infer_db.connectors import BaseConnector

class ExistingConstraintStrategy(PKStrategy):
    """
    1. Trust the Database:
    If the database explicitly says "This is a PK", we believe it.
    """
    def detect(self, table_name: str, connector: BaseConnector):
        logger.info(f"Checking existing constraints for {table_name}")
        all_constraints = connector.get_existing_constraints()
        
        pks = [
            row['column_name'] 
            for row in all_constraints 
            if row['table_name'] == table_name and row['constraint_type'] == 'PRIMARY KEY'
        ]
        return pks

class SurrogateKeyStrategy(PKStrategy):
    """
    2. The Sequence Check:
    Checks if a column behaves like a numeric sequence (1, 2, 3...).
    It must be UNIQUE and behave like an auto-incrementing integer.
    """
    def detect(self, table_name: str, connector: BaseConnector):
        logger.info(f"Checking surrogate keys for {table_name}")
        cols = connector.get_columns_metadata(table_name)
        row_count = connector.get_row_count(table_name)
        
        if row_count == 0:
            return []

        candidates = []
        
        for col in cols:
            col_name = col['column']
            default_val = str(col.get('column_default') or '')

            # A. Check by Definition (Explicit SERIAL/Identity)
            if 'nextval' in default_val.lower():
                candidates.append(col_name)
                continue 
            
            # B. Check by Heuristic (Math)
            # Only checking integers
            if 'int' not in col['dtype'].lower():
                continue

            stats = connector.get_column_stats(table_name, col_name)
            if not stats:
                continue

            # CRITICAL: A surrogate key MUST be unique.
            # This prevents FKs in junction tables (e.g. post_categories.post_id) 
            # from being detected as Surrogates just because they are integers.
            if stats['distinct'] != row_count:
                continue

            # Sequence Logic: (Max - Min + 1) == Total Rows?
            # E.g. Rows=3. Values=[1,2,3]. (3 - 1 + 1) = 3. Match!
            min_val = stats['min']
            max_val = stats['max']

            # Safety check for None or non-numeric types
            if isinstance(min_val, (int, float)) and isinstance(max_val, (int, float)):
                if (max_val - min_val + 1) == row_count:
                    candidates.append(col_name)

        return candidates

class DataValidationStrategy(PKStrategy):
    """
    3. The Reality Check (Type & Nulls):
    Rejects columns that are theoretically valid but bad practice (e.g. Float, JSON).
    Also checks REAL data for nulls.
    """
    def detect(self, table_name: str, connector: BaseConnector):
        logger.info(f"Checking data validation for {table_name}")
        cols = connector.get_columns_metadata(table_name)
        valid_candidates = []
        
        FORBIDDEN_TYPES = [
            'float', 'double', 'decimal', 'numeric', 'real',
            'bytea', 'json', 'jsonb', 'xml',
            'timestamp', 'time', 'interval', 'boolean'
        ]
        MAX_VARCHAR_LENGTH = 200

        for col in cols:
            col_name = col['column']
            dtype = col['dtype'].lower()
            char_len = col.get('char_len')

            # Rule A: Block forbidden types
            if any(t in dtype for t in FORBIDDEN_TYPES):
                continue
            
            # Rule B: Check Data Quality
            stats = connector.get_column_stats(table_name, col_name)
            if not stats:
                continue
            
            # THE REALITY CHECK:
            # Even if schema says 'nullable=YES', if data has NO nulls, it's valid.
            # If data HAS nulls, it cannot be a PK.
            if stats['nulls'] > 0:
                continue

            # Rule C: Length Check (allow UUIDs, reject long text)
            if 'uuid' in dtype:
                pass
            elif (char_len and char_len > MAX_VARCHAR_LENGTH):
                continue

            valid_candidates.append(col_name)
            
        return valid_candidates

class UniquenessStrategy(PKStrategy):
    """
    4. The Uniqueness Check:
    The most important test. If it identifies a row uniquely, it's a candidate.
    """
    def detect(self, table_name: str, connector: BaseConnector):
        logger.info(f"Checking uniqueness for {table_name}")
        metadata = connector.get_columns_metadata(table_name)
        total_rows = connector.get_row_count(table_name)
        
        if total_rows == 0:
            return []

        unique_cols = []
        for col in metadata:
            col_name = col['column']
            
            # Fetch stats (min, max, distinct, nulls)
            stats = connector.get_column_stats(table_name, col_name)
            
            if not stats:
                continue

            # 1. Null Check (Data Driven)
            # A PK cannot contain NULLs.
            if stats['nulls'] > 0:
                continue

            # 2. Uniqueness Check
            # If Count(Distinct Values) == Total Rows, it is a Unique Key.
            if stats['distinct'] == total_rows:
                unique_cols.append(col_name)
        
        return unique_cols