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 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,
                 endpoint: str = "druid/v2"):

        self.host = host
        self.port = port
        self.endpoint = endpoint
        self.client = PyDruid(
            f"http://{self.host}:{self.port}", {self.endpoint}
        )
        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_str = input_data['last_data_timestamp']
        if last_data_timestamp_str is None:
            last_data_timestamp = datetime(1970, 1, 1, 0, 0, 0)
        else:
            last_data_timestamp = datetime.strptime(
                last_data_timestamp_str, "%Y-%m-%d %H:%M:%S.%f")

        self.debug(
            f"Loading data from Druid: {input_data}", metadata=metadata)

        end_time = datetime(9999, 12, 31, 23, 59, 59)

        interval = f"{last_data_timestamp.isoformat()}Z/{end_time.isoformat()}Z"

        builder = QueryBuilder()

        query = builder.scan(
            {
                "datasource": datasource,
                "intervals": interval,
                "columns": ["timestamp", "value", "tag"],
                "limit": 10000,
            }
        )

        result = query.export_pandas()

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

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

        return result.to_dict(orient="records")