import concurrent.futures
import uuid
import psycopg2
import psycopg2.extras
from credativ_pg_migrator.anonymization.routing import MigratorAnonymizer

class AnonymizationWorkflow:
    def __init__(self, orchestrator):
        self.orchestrator = orchestrator
        self.config_parser = orchestrator.config_parser
        self.migrator_tables = orchestrator.migrator_tables
        self.logger = orchestrator.logger

        # Validate that both source and target are PostgreSQL
        if self.config_parser.get_source_db_type() != 'postgresql' or self.config_parser.get_target_db_type() != 'postgresql':
            raise ValueError("Anonymization workflow currently only supports PostgreSQL to PostgreSQL migrations.")

        self.anonymizer = MigratorAnonymizer(self.config_parser.config)

    def run(self):
        self.config_parser.print_log_message('INFO', "anonymization_workflow: run: Starting PG-to-PG Anonymization Workflow.")

        if not self.anonymizer.is_active():
            self.config_parser.print_log_message('WARNING', "anonymization_workflow: No anonymization rules found. Acting as standard PG-to-PG copy.")

        # Temporarily disable standard data migration to let standard table worker create tables without copying data
        original_migrate_data = self.config_parser.config['migration'].get('migrate_data', True)
        self.config_parser.config['migration']['migrate_data'] = False
        
        self.orchestrator.run_create_user_defined_types()
        self.orchestrator.run_create_domains()
        self.orchestrator.stdwf_migrate_sequences()

        self.orchestrator.stdwf_migrate_tables()

        # Restore original migrate_data flag
        self.config_parser.config['migration']['migrate_data'] = original_migrate_data

        # If data migration was requested, run our anonymization copy
        if original_migrate_data:
            self.anonymization_copy_data()

        self.orchestrator.stdwf_migrate_indexes('standard')
        self.orchestrator.stdwf_migrate_constraints()
        self.orchestrator.stdwf_migrate_views()
        self.orchestrator.stdwf_migrate_funcprocs()
        self.orchestrator.stdwf_migrate_triggers()
        self.orchestrator.stdwf_migrate_indexes('function_based')
        self.orchestrator.stdwf_migrate_comments()
        self.orchestrator.run_post_migration_script()

    def anonymization_copy_data(self):
        self.migrator_tables.insert_main({'task_name': 'Orchestrator', 'subtask_name': 'anonymization data copy'})
        workers_requested = self.config_parser.get_parallel_workers_count()

        self.config_parser.print_log_message('INFO', f"anonymization_workflow: Starting {workers_requested} parallel workers for anonymization copy.")

        raw_tables = self.migrator_tables.fetch_all_tables(only_unfinished=False)
        migrate_tables = []
        for table_row in raw_tables:
            table_data = self.migrator_tables.decode_table_row(table_row)
            source_rows_all = table_data.get('source_table_rows_all', 0)
            source_rows_limited = table_data.get('source_table_rows_limited', 0)
            
            # Simplified logic for example: just process tables with rows
            if source_rows_limited > 0:
                migrate_tables.append(table_data)
            else:
                self.migrator_tables.update_table_status({'row_id': table_data['id'], 'success': True, 'message': 'anonymized data OK (0 rows)'})

        if migrate_tables:
            with concurrent.futures.ThreadPoolExecutor(max_workers=workers_requested) as executor:
                futures = {}
                for table_data in migrate_tables:
                    while len(futures) >= workers_requested:
                        done, _ = concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_COMPLETED)
                        for future in done:
                            table_done = futures.pop(future)
                            if not future.result():
                                self.config_parser.print_log_message('ERROR', "anonymization_workflow: Stopping execution due to error.")
                                exit(1)
                            else:
                                self.migrator_tables.update_table_status({'row_id': table_done['id'], 'success': True, 'message': 'anonymized data OK'})

                    future = executor.submit(self.anonymization_data_worker, table_data)
                    futures[future] = table_data

                for future in concurrent.futures.as_completed(futures):
                    table_done = futures[future]
                    if not future.result():
                        self.config_parser.print_log_message('ERROR', "anonymization_workflow: Stopping execution due to error.")
                        exit(1)
                    else:
                        self.migrator_tables.update_table_status({'row_id': table_done['id'], 'success': True, 'message': 'anonymized data OK'})
        
        self.migrator_tables.update_main_status({'task_name': 'Orchestrator', 'subtask_name': 'anonymization data copy', 'success': True, 'message': 'finished OK'})

    def anonymization_data_worker(self, table_data):
        worker_id = uuid.uuid4()
        source_schema = table_data['source_schema_name']
        source_table = table_data['source_table_name']
        source_table_id = table_data.get('source_table_id', 0)
        source_table_rows_all = table_data.get('source_table_rows_all', 0)
        source_table_rows_limited = table_data.get('source_table_rows_limited', 0)
        target_schema = table_data['target_schema_name']
        target_table = table_data['target_table_name']
        target_table_rows = table_data.get('target_table_rows', 0)
        
        try:
            self.config_parser.print_log_message('INFO', f"anonymization_workflow: Worker {worker_id}: Processing {target_table}")

            protocol_id = self.migrator_tables.insert_data_migration({
                'worker_id': worker_id,
                'source_table_id': source_table_id,
                'source_schema_name': source_schema,
                'source_table_name': source_table,
                'target_schema_name': target_schema,
                'target_table_name': target_table,
                'source_table_rows_all': source_table_rows_all,
                'source_table_rows_limited': source_table_rows_limited,
                'target_table_rows': target_table_rows,
            })
            self.migrator_tables.update_data_migration_started(protocol_id)

            source_conn = self.orchestrator.load_connector('source')
            target_conn = self.orchestrator.load_connector('target')
            
            source_conn.connect()
            target_conn.connect()
            
            # Named cursors require a transaction (no autocommit)
            source_conn.connection.autocommit = False

            batch_size = self.config_parser.get_batch_size()
            
            target_columns = list(table_data.get('target_columns', {}).values())
            # Sort columns by column_id for correct ordering
            target_columns.sort(key=lambda x: x.get('column_id', 0))
            col_names = [col['column_name'] for col in target_columns]
            
            select_sql = f'SELECT {", ".join([f"{col}" for col in col_names])} FROM "{source_schema}"."{source_table}"'
            
            src_cursor = source_conn.connection.cursor(name=f"cursor_{worker_id}")
            src_cursor.execute(select_sql)

            import time
            total_inserted_rows = 0
            batch_number = 0
            batch_durations = []

            while True:
                batch_start_time = time.time()
                reading_start_time = batch_start_time
                rows = src_cursor.fetchmany(batch_size)
                if not rows:
                    break
                
                batch_number += 1
                reading_duration = time.time() - reading_start_time
                
                transforming_start_time = time.time()
                
                col_max_lengths = {}
                for col in target_columns:
                    max_len = col.get('character_maximum_length')
                    if max_len is not None and max_len != '':
                        try:
                            col_max_lengths[col['column_name']] = int(max_len)
                        except ValueError:
                            pass

                # Convert rows to dicts for anonymizer
                formatted_data = []
                for idx, row in enumerate(rows):
                    row_dict = dict(zip(col_names, row))
                    if self.anonymizer.is_active():
                        row_dict = self.anonymizer.anonymize_row(target_table, row_dict)
                    
                    # Truncate string values that exceed column limits
                    for col_name, max_len in col_max_lengths.items():
                        if isinstance(row_dict.get(col_name), str):
                            if len(row_dict[col_name]) > max_len:
                                row_dict[col_name] = row_dict[col_name][:max_len]
                                
                    formatted_data.append([row_dict[col] for col in col_names])

                    if idx > 0 and idx % 10000 == 0:
                        self.config_parser.print_log_message('INFO', f"anonymization_workflow: Worker {worker_id}: Anonymized {idx} rows of {target_table} in current batch...")
                
                transforming_duration = time.time() - transforming_start_time
                
                # Determine __RAW_SQL__ injection
                if formatted_data:
                    first_row = formatted_data[0]
                    new_insert_values_list = []
                    raw_sql_indices = set()
                    
                    for i, val in enumerate(first_row):
                        if isinstance(val, str) and val.startswith('__RAW_SQL__:'):
                            raw_sql_content = val.split('__RAW_SQL__:', 1)[1]
                            new_insert_values_list.append(raw_sql_content)
                            raw_sql_indices.add(i)
                        else:
                            new_insert_values_list.append('%s')
                            
                    insert_values_str = ', '.join(new_insert_values_list)
                    insert_query = f'INSERT INTO "{target_schema}"."{target_table}" ({", ".join(col_names)}) VALUES ({insert_values_str})'
                    
                    # Remove RAW SQL items from tuples
                    if raw_sql_indices:
                        cleaned_data = []
                        for row in formatted_data:
                            new_row = tuple(val for i, val in enumerate(row) if i not in raw_sql_indices)
                            cleaned_data.append(new_row)
                        formatted_data = cleaned_data
                    
                    # Execute batch
                    writing_start_time = time.time()
                    tgt_cursor = target_conn.connection.cursor()
                    psycopg2.extras.execute_batch(tgt_cursor, insert_query, formatted_data, page_size=batch_size)
                    target_conn.connection.commit()
                    tgt_cursor.close()
                    writing_duration = time.time() - writing_start_time
                    
                total_inserted_rows += len(formatted_data)
                batch_end_time = time.time()
                batch_duration = batch_end_time - batch_start_time
                batch_durations.append(batch_duration)

                import datetime
                batch_start_dt = datetime.datetime.fromtimestamp(batch_start_time)
                batch_end_dt = datetime.datetime.fromtimestamp(batch_end_time)
                
                self.migrator_tables.insert_batches_stats({
                    'source_schema_name': source_schema,
                    'source_table_name': source_table,
                    'source_table_id': source_table_id,
                    'chunk_number': 1,
                    'batch_number': batch_number,
                    'batch_start': batch_start_dt.strftime('%Y-%m-%d %H:%M:%S.%f'),
                    'batch_end': batch_end_dt.strftime('%Y-%m-%d %H:%M:%S.%f'),
                    'batch_rows': len(formatted_data),
                    'batch_seconds': batch_duration,
                    'worker_id': worker_id,
                    'reading_seconds': reading_duration,
                    'transforming_seconds': transforming_duration,
                    'writing_seconds': writing_duration if 'writing_duration' in locals() else 0,
                })

            src_cursor.close()
            source_conn.disconnect()
            target_conn.disconnect()
            
            shortest_batch_seconds = min(batch_durations) if batch_durations else 0
            longest_batch_seconds = max(batch_durations) if batch_durations else 0
            average_batch_seconds = sum(batch_durations) / len(batch_durations) if batch_durations else 0
            
            self.migrator_tables.update_data_migration_status({
                'row_id': protocol_id,
                'success': True,
                'message': 'anonymized data OK',
                'target_table_rows': total_inserted_rows,
                'batch_count': batch_number,
                'shortest_batch_seconds': shortest_batch_seconds,
                'longest_batch_seconds': longest_batch_seconds,
                'average_batch_seconds': average_batch_seconds,
            })
            
            return True

        except Exception as e:
            self.config_parser.print_log_message('ERROR', f"anonymization_workflow: Worker {worker_id}: Error processing {target_table} -> {e}")
            if 'protocol_id' in locals():
                self.migrator_tables.update_data_migration_status({
                    'row_id': protocol_id,
                    'success': False,
                    'message': str(e),
                    'target_table_rows': total_inserted_rows if 'total_inserted_rows' in locals() else 0,
                    'batch_count': batch_number if 'batch_number' in locals() else 0,
                    'shortest_batch_seconds': 0,
                    'longest_batch_seconds': 0,
                    'average_batch_seconds': 0,
                })
            return False
