from infer_db.utils.logger_util import logger
from infer_db.connectors import BaseConnector
from infer_db.strategies import PKStrategy, DataValidationStrategy, SurrogateKeyStrategy, FKStrategy

class InferenceEngine:
    def __init__(self, connector: BaseConnector):
        self.connector = connector
        self.pk_strategies: list[PKStrategy] = []
        self.fk_strategies: list[FKStrategy] = []

    def add_pk_strategy(self, strategy: PKStrategy):
        self.pk_strategies.append(strategy)

    def add_fk_strategy(self, strategy: FKStrategy):
        self.fk_strategies.append(strategy)

    def run(self):
        tables = self.connector.get_tables()
        schema_report = {}

        # =================================================
        # PHASE 1: PRIMARY KEY DETECTION (Build the Map)
        # =================================================
        logger.info("--- PHASE 1: Detecting Primary Keys ---")
        for table in tables:
            metadata = self.connector.get_columns_metadata(table)
            
            # We'll use a dict to track candidates and their 'tags'
            # e.g., {'id': {'is_surrogate': True}}
            candidates = {}

            for strategy in self.pk_strategies:
                # Skip the validator during the discovery phase
                if isinstance(strategy, DataValidationStrategy):
                    continue
                    
                found_cols = strategy.detect(table, self.connector)
                
                for col in found_cols:
                    if col not in candidates:
                        candidates[col] = {"is_surrogate": False}
                    
                    # If the Surrogate strategy found it, tag it!
                    if isinstance(strategy, SurrogateKeyStrategy):
                        candidates[col]["is_surrogate"] = True

            # Run Validation Filter
            validator = DataValidationStrategy()
            valid_cols = validator.detect(table, self.connector)
            
            # Final cleanup: keep only valid candidates
            final_pks = {
                col: attr for col, attr in candidates.items() 
                if col in valid_cols
            }            

            schema_report[table] = {
                "pks": final_pks, 
                "columns": metadata
            }
        
        # =================================================
        # PHASE 2: FOREIGN KEY DETECTION (Connect the Dots)
        # =================================================
        logger.info("--- PHASE 2: Detecting Foreign Keys ---")
        all_relations = []
        
        # We check every table against every other table
        for source_table in tables:
            source_metadata = schema_report[source_table]['columns']
            # Get Source PKs to prevent using them as FK candidates
            source_pks_list = list(schema_report[source_table]['pks'].keys())
            
            for target_table in tables:
                if source_table == target_table:
                    continue # Skip self-referencing for now (unless desired)

                # Prepare Target PK Info (Need Types for compatibility check)
                target_pks_refined = {}
                target_pk_list = schema_report[target_table]['pks']
                target_all_cols = schema_report[target_table]['columns']

                # Augment the simple PK list with full metadata (Type, Length, etc.)
                for pk_name, pk_attrs in target_pk_list.items():
                    full_col_def = next((c for c in target_all_cols if c['column'] == pk_name), None)
                    if full_col_def:
                        target_pks_refined[pk_name] = {
                            **pk_attrs, 
                            'dtype': full_col_def['dtype']
                        }

                # Run FK Strategies
                for strategy in self.fk_strategies:
                    # strategies now accept 'connector' to run queries/cache
                    fks = strategy.detect(
                        source_table=source_table, 
                        target_table=target_table, 
                        source_cols=source_metadata, 
                        target_pks=target_pks_refined,
                        connector=self.connector,
                        source_pks_list=source_pks_list
                    )
                    all_relations.extend(fks)

        # Attach relations to the final report
        schema_report['relations'] = all_relations

        return schema_report