# Copyright © 2024-2025 Pavel Rabaev
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
logger = logging.getLogger(__name__)

import os
import sys
import time
from abc import ABC, abstractmethod
from petaly.core.composer import Composer
from petaly.utils.utils import measure_time
from petaly.utils.file_handler import FileHandler
from petaly.core.object_metadata import ObjectMetadata
from petaly.core.type_mapping import TypeMapping
from petaly.core.data_object import DataObject


class DBExtractor(ABC):
    """Abstract base class for database extractors.
    
    This class provides the core functionality for extracting data from database sources.
    It handles metadata extraction, query composition, and data export to CSV files.
    
    Key responsibilities:
    - Extracts metadata from database objects
    - Composes and executes extraction queries
    - Manages file output and metadata storage
    - Handles type mapping and data transformation
    
    Attributes:
        pipeline: The pipeline instance containing configuration and state
        f_handler: FileHandler instance for file operations
        composer: Composer instance for query composition
        m_conf: Main configuration instance
        type_mapping: TypeMapping instance for data type conversions
        object_metadata: ObjectMetadata instance for metadata management
    """

    def __init__(self, pipeline):
        super().__init__()

        self.pipeline = pipeline
        self.f_handler = FileHandler()
        self.composer = Composer()
        self.m_conf = self.pipeline.m_conf
        self.type_mapping = TypeMapping(pipeline)
        self.object_metadata = ObjectMetadata(pipeline)

        if self.m_conf.set_extractor_paths(self.pipeline.source_connector_id):
            self.connector_extract_to_stmt_fpath = self.m_conf.connector_extract_to_stmt_fpath
            self.connector_metadata_sql_fpath = self.m_conf.connector_metadata_sql_fpath
            self.query_origin = self.f_handler.load_file(self.connector_metadata_sql_fpath)

    @abstractmethod
    def extract_to(self, extractor_obj_conf):
        pass

    @abstractmethod
    def get_query_result(self, meta_query):
        pass

    @abstractmethod
    def compose_extract_to_stmt(self, extract_to_stmt, extract_config) -> dict:
        pass

    @measure_time
    def extract_data(self):
        """Extracts data from the database source and exports it to CSV files.
        
        This method orchestrates the entire extraction process:
        1. Cleans up the pipeline output directory
        2. Composes and executes metadata queries
        3. Processes metadata for each object
        4. Extracts data for each object to CSV files
        
        The method handles timing and logging of the extraction process.
        """

        logger.info(f"[--- Extract from {self.pipeline.source_connector_id} ---]")
        start_total_time = time.time()
        # 1. Start with cleanup
        self.f_handler.cleanup_dir(self.pipeline.output_pipeline_dpath)

        # 2. compose_extract_scripts
        meta_query = self.compose_meta_query()

        # 3. get meta query result, expected as a dict
        meta_query_result = self.execute_meta_query(meta_query)

        # 4. save metadata and export scripts
        object_list = self.object_metadata.process_metadata(meta_query_result)

        # 5. run loop for each object
        for object_name in object_list:

            logger.info(f"Extract object: {object_name} started...")
            start_time = time.time()

            # 5. get all export scripts and store data into output directory
            extractor_obj_conf = self.get_extractor_obj_conf(object_name)

            # 6. run export data
            self.extract_to(extractor_obj_conf)

            end_time = time.time()
            logger.info(f"Extract object: {object_name} completed | time: {round(end_time - start_time, 2)}s")

        end_total_time = time.time()
        logger.info(f"Extract completed, duration: {round(end_total_time - start_total_time, 2)}s")

    def execute_meta_query(self, meta_query):
        """Executes a metadata query and returns the results.
        
        The method executes the provided SQL query for metadata extraction.
        If the query cannot be executed, it raises a SystemExit error.
        """
        logger.debug("Execute meta-query and create extract scripts")
        if meta_query is not None:
            query_result = self.get_query_result(meta_query)

        else:
            logger.error(
                f"Meta Query for pipeline {self.pipeline.pipeline_name} can not be executed. Review your configuration.")

            query_result = None
        return query_result

    def get_extractor_obj_conf(self, object_name) ->dict:
        """Gets the configuration for extracting a specific object.
        
        The method composes a complete configuration dictionary containing:
        - Object name and paths
        - Metadata directory
        - Extract queries
        - Object settings
        - Blob prefix for cloud storage
        - Extract statements
        """

        extractor_obj_conf = {'object_name': object_name}

        # 1. compose metadata directory
        output_metadata_object_dir = self.pipeline.output_object_metadata_dpath.format(object_name=object_name)
        extractor_obj_conf.update({'output_metadata_object_dir': output_metadata_object_dir})
        metadata_fpath = self.pipeline.output_object_metadata_fpath.format(object_name=object_name)

        # 2. get and compose metadata query
        table_metadata = self.f_handler.load_file_as_dict(metadata_fpath, 'json')
        extract_queries_dict = self.compose_extract_queries(table_metadata)
        extractor_obj_conf.update(extract_queries_dict)

        # 3. add object_default_settings
        #extractor_obj_conf.update({'object_settings': table_metadata.get('object_settings')})
        data_object = self.get_data_object(object_name)
        logger.debug(f"The object settings combined with default settings: {data_object.object_settings}")
        extractor_obj_conf.update({'object_settings': data_object.object_settings})

        # blob-prefix, used for storage in cloud services (e.g. Redshift (s3), Bigquery (GCS))
        blob_prefix = self.composer.compose_bucket_object_path(self.pipeline.source_attr.get('bucket_pipeline_prefix'),
                                                            self.pipeline.pipeline_name,
                                                            object_name)
        extractor_obj_conf.update({'blob_prefix': blob_prefix})

        # 4. load stmt_extract_to.txt and transform it in later stage
        extract_to_stmt = self.f_handler.load_file(self.connector_extract_to_stmt_fpath)
        extract_to_stmt = self.compose_extract_to_stmt(extract_to_stmt, extractor_obj_conf)
        extractor_obj_conf.update({'extract_to_stmt': extract_to_stmt})

        # 5.  save extract_to_stmt under output_extract_to_file_fpath
        output_extract_to_stmt_fpath = self.pipeline.output_extract_to_stmt_fpath.format(object_name=object_name)
        extractor_obj_conf.update({'extract_to_stmt_fpath': output_extract_to_stmt_fpath})
        self.f_handler.save_file(output_extract_to_stmt_fpath, extract_to_stmt)

        # 6. compose output data dir
        output_data_object_dir = self.pipeline.output_object_data_dpath.format(object_name=object_name)
        extractor_obj_conf.update({'output_data_object_dir': output_data_object_dir})
        self.f_handler.make_dirs(output_data_object_dir)

        # compose output object file path
        output_object_fpath = os.path.join(output_data_object_dir, object_name + '.csv')
        extractor_obj_conf.update({'output_object_fpath': output_object_fpath})

        logger.debug(f"Config for data extract: {extractor_obj_conf}")
        return extractor_obj_conf

    def compose_meta_query(self):
        """Composes the metadata query based on pipeline configuration.
        
        The query is built using:
        - Schema information
        - Table list (if specified)
        - Column definitions
        """
        def get_table_stmt(data_objects_list):
            if len(data_objects_list)>0:
                table_stmt = 'AND tb.table_name IN ({tbl_list})'
                table_string = ''
                for tbl in data_objects_list:
                    table_string += "'" + tbl + "',"
                         
                table_string = table_string.rstrip(',')
                return table_stmt.format(tbl_list=table_string)
            else:
                return ''

        logger.debug("Compose data source meta query:")
        
        # if data_objects_from_cli is set, use it to compose the table_stmt and ignore all other settings
        if len(self.pipeline.data_objects_from_cli)>0:
            table_stmt = get_table_stmt(self.pipeline.data_objects_from_cli)
        
        elif self.pipeline.data_attributes.get('data_objects_spec_mode') in ("ignore","prefer"):
            table_stmt = ''
            
        else:
        # it means data_objects_spec_mode is set to "only" and data_objects_spec: [] should has at least one object specification
        
            if len(self.pipeline.data_objects)==0:
                logger.warning(f"Pipeline {self.pipeline.pipeline_name} in {self.pipeline.pipeline_fpath} wasn't specified properly. If data_objects_spec_mode is set to \"only\" the data_objects_spec: [] should has at least one object specification")
                sys.exit()
			
            table_stmt = get_table_stmt(self.pipeline.data_objects)

        source_schema = self.pipeline.source_attr.get('database_schema')

        if source_schema is None:
            if self.f_handler.check_dict_key_exist(self.pipeline.source_attr, 'database_schema'):
                logger.warning(f"A source database schema wasn't specified. To continue, specify database_schema in pypeline.yaml ")
                sys.exit()

            source_schema = self.pipeline.source_attr.get('database_name')

        meta_query = self.query_origin.format(schema=source_schema, table_statement_list=table_stmt)

        logger.debug(f"Meta Query:\n {meta_query}")

        return meta_query

    @measure_time
    def compose_extract_queries(self, dict_obj):
        """Composes extraction queries for a specific object.
        
        The method processes object metadata to create:
        - Source schema and object names
        - Column list with type transformations
        """
        extract_obj_conf = {}

        if dict_obj.get('source_object_name') is not None:
            column_list = ''

            transformation = self.type_mapping.get_extractor_type_transformer()
            col_obj = dict_obj.get('columns')

            for idx, val in enumerate(col_obj):

                column_transformation = transformation.get(val['data_type'])

                if column_transformation is not None:
                    column_list += column_transformation.format(column_name=val['column_name']) + ","
                else:
                    column_list += f" {self.db_connector.metaquery_quote}{val['column_name']}{self.db_connector.metaquery_quote},"
                    #columns += self.normalise_column_name(val['column_name']) + ","

            column_list = column_list.rstrip(',')

            extract_obj_conf.update({
                                    'source_schema_name':dict_obj['source_schema_name'],
                                    'source_object_name':dict_obj['source_object_name'],
                                    'column_list':column_list})

            return extract_obj_conf

    def get_data_object(self, object_name):
        """Gets a DataObject instance for the specified object.
        
        Creates and returns a DataObject instance containing the object's
        configuration and settings.
        """
        return DataObject(self.pipeline, object_name)
