from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    import pandas as pd
    from typing import List, Optional, Any
    from datetime import datetime, timedelta
    from pydruid.client import PyDruid
    from pydruid.query import QueryBuilder
    from sqlalchemy.engine import create_engine
    from sqlalchemy import MetaData, Table, select, text
    from sientia_do.temporal.activities.base import BaseActivity
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.temporal.utils.logger import Logger


class Druid(BaseActivity):
    def __init__(self, host: str, port: int,
                 logger: Logger, notification_handler: NotificationHandler):

        self.host = host
        self.port = port
        self.engine = create_engine(
            f'druid://{self.host}:{self.port}/druid/v2/sql/')
        logger.info(
            f"Druid client initialized with host: {self.host}, port: {self.port}")

        BaseActivity.__init__(self, logger=logger,
                              notification_handler=notification_handler)

    def shutdown(self):
        self.client.close()

    def __del__(self):
        self.shutdown()

    @activity.defn(name="load_latest_druid_data")
    async def load_latest_druid_data(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Loads the latest data from Druid.
        """
        metadata = input_data['metadata']
        datasource = f"raw_{input_data['schedule_name']}"
        last_data_timestamp = input_data['last_data_timestamp']
        self.debug(
            f"Loading data from Druid: {input_data}", metadata=metadata)

        query = f'"__time" > TIMESTAMP \'{last_data_timestamp}\''

        self.info(
            f"Loading data from Druid: {datasource} with query: {query}"
        )

        places = Table(datasource, MetaData(), autoload_with=self.engine)
        stmt = select(places).where(text(query))

        result = pd.read_sql(stmt, self.engine)

        result["inserted_at"] = pd.to_datetime(result["__time"]).dt.strftime(
            "%Y-%m-%d %H:%M:%S.%f")

        result.drop(columns=["__time"], inplace=True)

        self.info(
            f"Loaded {len(result)} rows from Druid"
        )

        self.debug(
            f"Druid query result: {result}", metadata=metadata)

        return result.to_dict()