from ..snowflake.base import SnowflakeJob
from datetime import datetime, timezone
from src.utils.database import write_to_snowflake
from src.utils.snowflake.constants import QueryType, Role, Warehouse
from .constants import JobGroup
from dagster import OpExecutionContext
from typing import Optional
import pandas as pd
import os
import requests
import asyncio
import aiohttp
import time


class ClerkRefundCandidates(SnowflakeJob):
    def __init__(self):
        super().__init__(
            name="clerk_refund_candidates",
            description="Applies bot metadata to Clerk users so we can get refunds on bots",
            schedule="*/10 * * * *",
            query_file="../bots/queries/clerk_refund_candidates.sql",
            query_type=QueryType.SELECT,
            warehouse=Warehouse.SMALL,
            role=Role.ACCOUNTADMIN,
            group_name="finops",
            # monitored=True,
            owners=["team:core-pod"],
            metadata={
                "slack": "#tech-anti-bots",
            },
            tags={"team": "core-pod", "category": "finops", "tech-alerts": "true"},
            deps=["bot_status_changes"],
        )

    async def update_clerk_metadata(self, session, user_id, headers):
        """Async function to update Clerk metadata for a single user."""
        try:
            payload = {"private_metadata": {"is_bot": "true"}}

            async with session.patch(
                f"https://api.clerk.com/v1/users/{user_id}/metadata", headers=headers, json=payload
            ) as response:
                response_text = await response.text()
                return {
                    "CLERK_ID": user_id,
                    "STATUS_CODE": response.status,
                    "ERROR": response_text if response.status != 200 else None,
                    "CREATED_AT": datetime.now(tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
                    "DAGSTER_RUN_ID": self.context.run_id,
                }
        except Exception as e:
            return {
                "CLERK_ID": user_id,
                "STATUS_CODE": None,
                "ERROR": str(e),
                "CREATED_AT": datetime.now(tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
                "DAGSTER_RUN_ID": self.context.run_id,
            }

    async def process_batch(self, batch, headers):
        """Process a batch of users concurrently."""
        async with aiohttp.ClientSession() as session:
            tasks = [
                self.update_clerk_metadata(session, row["CLERK_ID"], headers)
                for _, row in batch.iterrows()
            ]
            return await asyncio.gather(*tasks)

    def post_execute(self, result: Optional[pd.DataFrame], context: OpExecutionContext) -> None:
        """Update Clerk metadata for each user and track results."""
        if result is None or len(result) == 0:
            context.log.info("No users to update in Clerk")
            return

        # Convert result to DataFrame if it's not already
        if not isinstance(result, pd.DataFrame):
            result = pd.DataFrame(result)

        # Get Clerk secret key from environment
        clerk_secret_key = os.getenv("CLERK_API_KEY")
        if not clerk_secret_key:
            raise ValueError("CLERK_API_KEY environment variable not set")

        # Store context and secret key for async functions
        self.context = context
        self.clerk_secret_key = clerk_secret_key

        context.log.info(f"Preview of Clerk Refund Candidates:\n{result.head(10).to_markdown()}")

        # Prepare headers for all requests
        headers = {
            "Authorization": f"Bearer {self.clerk_secret_key}",
            "Content-type": "application/json",
        }

        # Process in batches of 1000 (Clerk's rate limit is supposed to be 3K but we're still hitting 429 errors)
        batch_size = 1000
        all_results = []
        total_processed = 0
        success_count = 0
        failure_count = 0

        # Process batches
        for i in range(0, len(result), batch_size):
            batch = result.iloc[i : i + batch_size].copy()  # Add .copy() to prevent modification
            context.log.info(
                f"Processing batch {i // batch_size + 1} of {(len(result) + batch_size - 1) // batch_size}"
            )

            # Process batch concurrently
            batch_results = asyncio.run(self.process_batch(batch, headers))
            all_results.extend(batch_results)

            # Update counts
            for batch_result in batch_results:  # Renamed to avoid shadowing
                if batch_result["STATUS_CODE"] == 200:
                    success_count += 1
                else:
                    failure_count += 1

            total_processed += len(batch)

            # Write to Snowflake every 1000 records
            if len(all_results) >= 1000:
                snowflake_df = pd.DataFrame(all_results)
                context.log.info(f"Writing batch of {len(snowflake_df)} records to Snowflake")
                context.log.info(f"Snowflake DataFrame:\n{snowflake_df.head(10).to_markdown()}")
                write_to_snowflake(snowflake_df, "DAGSTER_CLERK_REFUNDS")
                all_results = []  # Clear the list after writing

            # Wait 10 seconds between batches to respect rate limit
            if i + batch_size < len(result):
                context.log.info("Waiting 5 seconds before next batch...")
                time.sleep(5)

        # Write any remaining records
        if all_results:
            snowflake_df = pd.DataFrame(all_results)
            context.log.info(f"Writing final batch of {len(snowflake_df)} records to Snowflake")
            context.log.info(f"Snowflake DataFrame:\n{snowflake_df.head(10).to_markdown()}")
            write_to_snowflake(snowflake_df, "DAGSTER_CLERK_REFUNDS")

        # Log summary
        context.log.info(
            f"Clerk metadata update complete. "
            f"Total processed: {total_processed}, "
            f"Success: {success_count}, "
            f"Failures: {failure_count}"
        )


bots_hourly = ClerkRefundCandidates()