from .base import PostgresJob
from src.utils.monitoring import datadog_client
import pandas as pd
from typing import Optional, List
import time
from dagster import OpExecutionContext
from src.utils.database import write_to_snowflake

SNOWFLAKE_RESULTS_TABLE = "LABEL_MAKER_CREDITS_PROCESSED"


class LabelMakerCredits(PostgresJob):
    def __init__(self):
        super().__init__(
            name="label_maker_credits",
            description="Issues credits for Label Maker annotations",
            schedule="*/5 * * * *",  # runs every 5 minutes
            query_file="queries/label_maker_credits.sql",
            is_reader=False,
            group_name="ml_label_maker_credits",
            monitored=True,
            owners=["team:core-pod"],
            metadata={
                "slack": "#project-label-maker",
            },
            tags={"team": "core-pod", "category": "ml"},
        )

    def post_execute(
        self, data: Optional[List], context: OpExecutionContext
    ) -> Optional[pd.DataFrame]:
        """Convert the raw data into a DataFrame with proper columns"""
        if data is None:
            return None

        df = pd.DataFrame(
            data,
            columns=[
                "EMAIL",
                "USER_ID",
                "ANNOTATION_ID",
                "CREDITS_ISSUED",
                "CREDITS_PROCESSED",
                "CURRENT_BALANCE",
            ],
        )

        # Add dagster_run_id
        df["DAGSTER_RUN_ID"] = context.run_id

        # Preview the DataFrame
        context.log.info(f"\nDataFrame Preview:\n{df.head(10).to_markdown()}")
        context.log.info(
            f"\nDataFrame Summary:\n"
            + f"Total Rows: {len(df)}\n"
            + f"Total Credits Issued: {df['CREDITS_ISSUED'].sum()}\n"
            + f"Unique Users: {df['USER_ID'].nunique()}"
        )

        # Send metrics to Datadog
        now = int(time.time())
        metrics = [
            {
                "metric": "label_maker.credits.issued",
                "points": [(now, float(df["CREDITS_ISSUED"].sum()))],
                "type": "gauge",
                "tags": ["source:label_maker", f"run_id:{context.run_id}"],
            },
            {
                "metric": "label_maker.users.credited",
                "points": [(now, float(df["USER_ID"].nunique()))],
                "type": "gauge",
                "tags": ["source:label_maker", f"run_id:{context.run_id}"],
            },
            {
                "metric": "label_maker.annotations.processed",
                "points": [(now, float(len(df)))],
                "type": "gauge",
                "tags": ["source:label_maker", f"run_id:{context.run_id}"],
            },
        ]
        datadog_client.Metric.send(metrics=metrics)

        # Write to Snowflake
        write_to_snowflake(df, SNOWFLAKE_RESULTS_TABLE)
        return df


label_maker_credits = LabelMakerCredits()