# coding: utf-8
import os
from typing import Iterator, Any
import logging

import pandas as pd
from datetime import datetime, timezone

import dlt
from ruamel.yaml import YAML
from dlt.sources import TDataItems
from dlt.sources.filesystem import FileItemDict, filesystem, read_csv
from dlt.common.schema.typing import TSimpleRegex
#from dlt.destinations.adapters import bigquery_adapter

dlt.config["extract.workers"] = "16"
dlt.config["normalize.workers"] = "8"
dlt.config["runtime.log_level"] = "INFO"
dlt.config["runtime.dlthub_telemetry"] = "false"

# dlt.config["sources.data_writer.buffer_max_items"] = "100"
# dlt.config["sources.data_writer.file_max_items"] = "100000"
# dlt.config["sources.data_writer.file_max_bytes"] = "1000000"

dlt.config["destination.bigquery.location"] = "us-west2"
dlt.config["destination.bigquery.partition_expiration_days"] = "30"

@dlt.transformer(parallelized=True)
def read_csv_custom(
    items: Iterator[FileItemDict], chunksize: int = 10000, latest_timestamp: datetime = None, **pandas_kwargs: Any
) -> Iterator[TDataItems]:
    """Reads csv file with Pandas chunk by chunk.

    Args:
        chunksize (int): Number of records to read in one chunk
        **pandas_kwargs: Additional keyword arguments passed to Pandas.read_csv
    Returns:
        TDataItem: The file content
    """

    # apply defaults to pandas kwargs
    kwargs = {**{"header": "infer", "chunksize": chunksize}, **pandas_kwargs}

    for file_obj in items:
        # Here we use pandas chunksize to read the file in chunks and avoid loading the whole file
        # in memory.
        blob_timestamp = file_obj.get('file_name', None)
        if blob_timestamp is not None and '-' in blob_timestamp:
            blob_timestamp = blob_timestamp.split('-')[1]
            
            if blob_timestamp != latest_timestamp:
                logging.debug(f"Skipping {blob_timestamp} since it is not the latest timestamp {latest_timestamp}")
                continue

            blob_timestamp = datetime.fromtimestamp(int(int(latest_timestamp) / 1000))

            # try to read the file with the default encoding
            try:
                file = file_obj.open()
                df = pd.read_csv(file, **kwargs)
            except UnicodeDecodeError:
                file = file_obj.open()
                df = pd.read_csv(file, **kwargs, encoding='iso-8859-1')

            for df in pd.read_csv(file, **kwargs):
                df['source_created_timestamp'] = blob_timestamp
                yield df.to_dict(orient="records")

@dlt.transformer(parallelized=True)
def read_parquet_custom(
    items: Iterator[FileItemDict],
    chunksize: int = 10,
    latest_timestamp: datetime = None,
) -> Iterator[TDataItems]:
    from pyarrow import parquet as pq

    for file_obj in items:
        blob_timestamp = file_obj.get('file_name', None)
        if blob_timestamp is not None and '-' in blob_timestamp:
            blob_timestamp = blob_timestamp.split('-')[1]
            
            if blob_timestamp != latest_timestamp:
                logging.debug(f"Skipping {blob_timestamp} since it is not the latest timestamp {latest_timestamp}")
                continue
            blob_timestamp = datetime.fromtimestamp(int(int(latest_timestamp) / 1000))

            with file_obj.open() as f:
                parquet_file = pq.ParquetFile(f)
                for rows in parquet_file.iter_batches(batch_size=chunksize):
                    # Potentially avoid using pandas here
                    df = rows.to_pandas()
                    df['source_created_timestamp'] = blob_timestamp
                    yield df.to_dict(orient="records")


@dlt.source()
def gcs_source(source_gcs_bucket, source_tables, file_format,credentials=dlt.secrets.value):

    resources = []
    for table in source_tables:
        source_gcs_glob = table['gcs_glob']
        target_bq_table = table.get('identifier', table['name'])
        primary_keys    = table.get('primary_keys', ['id'])

        filesystem_resource = filesystem(
            bucket_url=source_gcs_bucket,
            file_glob=source_gcs_glob,
            credentials=credentials,
        )
        filesystem_resource.apply_hints(incremental=dlt.sources.incremental("modification_date", initial_value=datetime(2025,3,9,0,0,0,0,timezone.utc)))
        
        latest_timestamp = max([
            x['file_name'].split('-')[1] if (x.get('file_name', None) is not None and '-' in x.get('file_name', None))
            else 0 for x in filesystem_resource
        ])
        logging.info(f"Latest timestamp for {table}: {latest_timestamp}")

        if file_format == 'parquet':
            filesystem_pipe = filesystem_resource | read_parquet_custom(latest_timestamp=latest_timestamp)
        elif file_format == 'csv':
            filesystem_pipe = filesystem_resource | read_csv_custom(latest_timestamp=latest_timestamp)
        elif file_format == 'jsonl':
            filesystem_pipe = filesystem_resource | dlt.read_json()
        
        filesystem_pipe.parallelize()
        
        resources.append(filesystem_pipe.with_name(target_bq_table))
        # resources.append(
        #     bigquery_adapter(
        #         data=filesystem_pipe.with_name(target_bq_table),
        #         partition="source_created_timestamp",
        #         partition_expiration_days=30
        #     )
        # )
    
    return resources

import functions_framework

@functions_framework.cloud_event
def hello_pubsub(cloud_event):
    yaml = YAML()

    pipeline_dir = os.path.dirname(__file__)
    sources_path = os.path.join(pipeline_dir, "sources.yml")
    with open(sources_path, "r") as f:
        sources = yaml.load(f)['sources']

    for source in sources:
        if not source.get('enabled', True):
            logging.info(f"Skipping source {source['name']} since it is not enabled")
            continue
        source_name = source['name']
        source_format = source.get('file_format', 'csv')
        source_column_hints = source.get('column_hints', {})
        source_column_config = {x['name']: x['columns'] for x in source.get('tables', []) if x.get('columns', None) is not None}
        
        if source_format not in ['csv', 'parquet']:
            raise ValueError(f"Unsupported file format {source_format}")
        
        source_gcs_project = source['source_project']
        source_gcs_bucket = source['source_bucket']
        
        target_bq_project = source['database']
        target_bq_dataset = source['schema']

        pipeline = dlt.pipeline(
            pipeline_name=f"gcs_{source_format}_to_bq_{source_name}",
            destination='bigquery',
            dataset_name=target_bq_dataset,
            export_schema_path="schemas/export",
            progress="log",
        )

        source = gcs_source.clone(name=f"{source_format}_source")(source_gcs_bucket, source_tables=source['tables'], file_format=source_format)
        source.schema.add_type_detection("timestamp")
        source.schema.update_preferred_types(
            {
                "re:timestamp": "timestamp",
                "re:id": "text", # Tenant id can be larger than INT64
                "re:price": "text", # To accomodate for currency symbols / high-precision decimals
            } | source_column_hints
        )

        # for table in source.resources:
        #     source.resources[table].apply_hints(columns=source_column_config[table])

        info = pipeline.run(
            data=source,
            write_disposition={"disposition": "replace"}
        )
        print(info)


if __name__ == "__main__":
    hello_pubsub(None)