import sqlite3
import psycopg2
from django.core.management.base import BaseCommand
from django.conf import settings
from psycopg2 import sql

class Command(BaseCommand):
    help = 'Migrate data from SQLite to PostgreSQL'

    def get_column_types(self, sqlite_cursor, table):
        sqlite_cursor.execute(f"PRAGMA table_info({table})")
        columns_info = sqlite_cursor.fetchall()
        column_types = {col[1]: col[2] for col in columns_info}
        return column_types

    def get_indexes(self, sqlite_cursor, table):
        sqlite_cursor.execute(f"PRAGMA index_list({table})")
        indexes = sqlite_cursor.fetchall()
        index_details = []
        for index in indexes:
            index_name = index[1]
            index_info = {
                'name': index_name,
                'columns': self.get_index_columns(sqlite_cursor, index_name)
            }
            index_details.append(index_info)
        return index_details

    def get_index_columns(self, sqlite_cursor, index_name):
        sqlite_cursor.execute(f"PRAGMA index_info({index_name})")
        index_info = sqlite_cursor.fetchall()
        columns = [info[3] for info in index_info]
        return columns

    def map_sqlite_to_postgres_type(self, sqlite_type):
        # Define PostgreSQL column types (basic mapping example)
        type_mapping = {
            'TEXT': 'TEXT',
            'INTEGER': 'INTEGER',
            'REAL': 'REAL',
            'BLOB': 'BYTEA',
            'datetime': 'TIMESTAMP',
            'INT': 'INTEGER'
        }
        # Default to TEXT if SQLite type is not found in mapping
        return type_mapping.get(sqlite_type.upper(), 'TEXT')

    def create_table_if_not_exists(self, pg_cursor, table, columns_info):
        pg_column_definitions = []
        primary_key_columns = []

        for column in columns_info:
            # Access column information
            cid, name, col_type, notnull, dflt_value, pk = column
            pg_type = self.map_sqlite_to_postgres_type(col_type)
            notnull_clause = 'NOT NULL' if notnull else ''
            default_clause = f"DEFAULT '{dflt_value}'" if dflt_value else ''

            # Check if the column is a primary key
            if pk:
                # Use SERIAL type for primary key columns
                pg_type = 'SERIAL' if pg_type == 'INTEGER' else pg_type
                primary_key_columns.append(name)
            
            column_definition = f"{name} {pg_type} {notnull_clause} {default_clause}".strip()
            pg_column_definitions.append(column_definition)

        columns_definition_str = ', '.join(pg_column_definitions)
        primary_key_str = f", PRIMARY KEY ({', '.join(primary_key_columns)})" if primary_key_columns else ''
        
        # Create table query
        pg_create_table_query = f"""
            CREATE TABLE IF NOT EXISTS {table} (
                {columns_definition_str}{primary_key_str}
            );
        """
        
        # Define index queries
        index_queries = []
        indexes = self.get_indexes(pg_cursor, table)
        for index in indexes:
            index_name = index['name']
            index_columns = ', '.join(index['columns'])
            index_query = f"""
                CREATE INDEX IF NOT EXISTS {index_name}_idx
                ON {table} ({index_columns});
            """
            index_queries.append(index_query)

        try:
            # Execute the table creation query
            pg_cursor.execute(pg_create_table_query)
            self.stdout.write(self.style.SUCCESS(f"Table {table} created or already exists."))
            
            # Execute the index queries
            for query in index_queries:
                pg_cursor.execute(query)
                self.stdout.write(self.style.SUCCESS(f"Index {query.split()[2]} created or already exists."))

            pg_cursor.connection.commit()
        except psycopg2.Error as e:
            self.stdout.write(self.style.ERROR(f"Error creating table {table} or indexes: {e}"))
            pg_cursor.connection.rollback()

    def handle(self, *args, **kwargs):
        # Connect to SQLite database
        sqlite_conn = sqlite3.connect('/var/www/html/python/core_market_data.db',
            detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES)
        sqlite_cursor = sqlite_conn.cursor()

        # Connect to PostgreSQL database using settings
        pg_conn = psycopg2.connect(
            dbname='core_market_data_new',
            user=settings.DATABASES['default']['USER'],
            password=settings.DATABASES['default']['PASSWORD'],
            host=settings.DATABASES['default']['HOST'],
            port=settings.DATABASES['default']['PORT']
        )
        pg_cursor = pg_conn.cursor()

        # List of tables to migrate
        tables = ['company_balance_sheets_annual']  # Add your tables here

        for table in tables:
            try:
                # Fetch and print column details from SQLite
                sqlite_cursor.execute(f"PRAGMA table_info({table})")
                columns_info = sqlite_cursor.fetchall()
                  # Print column details with headings
                self.stdout.write(self.style.SUCCESS(f"Table: {table}"))
                self.stdout.write(self.style.SUCCESS("Column Details:"))
                self.stdout.write(self.style.SUCCESS(f"{'ID':<5} {'Name':<30} {'Type':<15} {'Not Null':<10} {'Default':<20} {'PK':<5}"))
                self.stdout.write(self.style.SUCCESS('-' * 80))
                
                for column in columns_info:
                    cid, name, col_type, notnull, dflt_value, pk = column
                    notnull = 'YES' if notnull else 'NO'
                    pk = 'YES' if pk else 'NO'
                    # Handle None values for dflt_value
                    dflt_value = dflt_value if dflt_value is not None else 'N/A'
                    self.stdout.write(self.style.SUCCESS(f"{cid:<5} {name:<30} {col_type:<15} {notnull:<10} {dflt_value:<20} {pk:<5}"))
                
                self.stdout.write(self.style.SUCCESS(""))
            except Exception as e:
                self.stdout.write(self.style.ERROR(f"Error migrating table {table}: {e}"))

        # Close connections
        sqlite_conn.close()
        pg_conn.close()

        self.stdout.write(self.style.SUCCESS('Successfully retrieved column and index details from SQLite and created tables in PostgreSQL'))