from datetime import datetime, timedelta, timezone
from typing import Dict, List, Optional
import json
import os
import requests

import dagster as dg
from dagster import AssetKey, DagsterInstance

# Slack configuration from environment variables
SLACK_TOKEN = os.environ.get("SLACK_TOKEN", "")
SLACK_CHANNEL = os.environ.get("SLACK_CHANNEL", "data-alerts")

# SLA configuration
SLA_DEADLINE_HOUR = 12

# Asset filter configuration
# BLACKLIST: Exclude these assets from monitoring (exact match or prefix)
# Leave empty [] to monitor all assets
ASSET_BLACKLIST = [
    # Examples:
    # "test_asset",            # Exclude specific asset
    # "temp_",                 # Exclude all assets starting with "temp_"
    # "user_top_score_songs",  # Exclude this specific asset
]


def _is_asset_monitored(asset_name: str, blacklist: List[str]) -> bool:
    """
    Check if an asset should be monitored based on blacklist.
    
    Args:
        asset_name: Name of the asset to check
        blacklist: List of asset names or prefixes to exclude
    
    Returns:
        True if asset should be monitored, False if blacklisted
    """
    # Check if asset matches any blacklist pattern
    for pattern in blacklist:
        if asset_name == pattern or asset_name.startswith(pattern):
            return False
    
    # Not in blacklist, monitor it
    return True


def _discover_daily_partitioned_assets(context: dg.SensorEvaluationContext, target_date, logger, blacklist=None):
    """
    Auto-discover daily partitioned assets from Dagster definitions with blacklist filtering.
    
    Args:
        context: Sensor evaluation context
        target_date: Target date to check
        logger: Logger instance
        blacklist: List of asset names/prefixes to exclude (None = use ASSET_BLACKLIST)
    """
    if blacklist is None:
        blacklist = ASSET_BLACKLIST
    
    daily_assets = []
    filtered_assets = []
    all_assets = context.repository_def.assets_defs_by_key
    
    logger.info(f"🔍 Auto-discovery: Starting daily asset discovery")
    if blacklist:
        logger.info(f"🚫 Blacklist enabled: {blacklist}")
        
    for asset_key, asset_def in all_assets.items():
        if hasattr(asset_def, 'partitions_def') and asset_def.partitions_def:
            # Check if it's a daily partitions definition
            if isinstance(asset_def.partitions_def, dg.DailyPartitionsDefinition):
                asset_name = asset_key.to_user_string()
                
                # Check if asset should be monitored
                if not _is_asset_monitored(asset_name, blacklist):
                    filtered_assets.append(asset_name)
                    logger.debug(f"⏭️  Filtered out: {asset_name}")
                    continue
                
                # Get start date from partition definition
                start_datetime = asset_def.partitions_def.start
                start_date = start_datetime.date() if hasattr(start_datetime, 'date') else start_datetime
                
                # Skip if target date is before asset start date
                if target_date < start_date:
                    logger.debug(f"Skipping {asset_name}: target date {target_date} before start date {start_date}")
                    continue
                
                daily_assets.append({
                    "asset_key": asset_key,
                    "asset_name": asset_name,
                    "expected_schedule": f"Daily (auto-discovered)",
                    "start_date": start_date.strftime("%Y-%m-%d")
                })
                logger.info(f"✅ Auto-discovered daily asset: {asset_name} (start: {start_date})")
    
    logger.info(f"🔍 Auto-discovery: Found {len(daily_assets)} monitored assets (filtered out {len(filtered_assets)})")
    return daily_assets


@dg.sensor(
    name="daily_jobs_sla_monitor",
    description="Automatically discovers and monitors ALL daily partitioned assets for missing partitions",
    minimum_interval_seconds=3600 * 24,  # Check every 24 hours
    metadata={
        "expected_data_cadence": "daily",
        "sla_deadline_hour": SLA_DEADLINE_HOUR,
        "slack_channel": SLACK_CHANNEL,
        "discovery_method": "automatic_with_fallback",
    }
)
def daily_jobs_sla_monitor(context: dg.SensorEvaluationContext) -> dg.SensorResult:
    """
    General sensor that monitors ALL daily partitioned assets for missing partitions.
    Sends a consolidated Slack alert if any daily jobs didn't run for yesterday's partition.    
    """
    logger = context.log
    current_utc = datetime.now(timezone.utc)
    target_date = current_utc.date() - timedelta(days=1)
    target_partition_key = target_date.strftime("%Y-%m-%d")
    
    logger.info(f"Checking daily jobs for partition: {target_partition_key}")
    
    try:
        instance = context.instance
        
        # Auto-discover daily partitioned assets from Dagster
        daily_assets = _discover_daily_partitioned_assets(context, target_date, logger)
        
        logger.info(f"Monitoring {len(daily_assets)} daily partitioned assets")
        
        # Check materialization status for each asset
        missing_assets = []
        checked_assets = []
        
        for asset_info in daily_assets:
            asset_key = asset_info["asset_key"]
            asset_name = asset_info["asset_name"]
            
            checked_assets.append(asset_name)
            
            # Check if this asset has materialization for the target partition
            has_materialization = _check_asset_materialization(
                instance, asset_key, asset_name, target_partition_key, logger
            )
            
            if not has_materialization:
                missing_assets.append({
                    "asset_name": asset_name,
                    "asset_key": asset_key,
                    "partition_key": target_partition_key,
                    "expected_schedule": asset_info["expected_schedule"],
                    "status": "missing_materialization"
                })
        
        logger.info(f"Checked {len(checked_assets)} assets, found {len(missing_assets)} missing partitions")
        
        if missing_assets:
            # Send consolidated alert for all missing partitions
            alert_message = _create_consolidated_alert_message(
                missing_assets, target_partition_key, current_utc
            )
            success = _send_slack_alert(alert_message, logger)
            success = False  # Set default since Slack alert is disabled
            
            logger.error(f"Missing partitions found: {len(missing_assets)} assets")
            return dg.SensorResult(
                cursor_result=target_partition_key,
                metadata={
                    "partition_checked": target_partition_key,
                    "total_assets_checked": len(daily_assets),
                    "missing_assets_count": len(missing_assets),
                    "missing_assets": [asset["asset_name"] for asset in missing_assets],
                    "alert_sent": success,
                    "check_time": current_utc.isoformat()
                }
            )
        else:
            # All assets have their partitions
            logger.info(f"All {len(daily_assets)} daily assets have partitions for {target_partition_key}")
            return dg.SensorResult(
                skip_reason=f"All daily assets have partitions for {target_partition_key}",
                cursor_result=target_partition_key,
                metadata={
                    "partition_checked": target_partition_key,
                    "total_assets_checked": len(daily_assets),
                    "missing_assets_count": 0,
                    "check_time": current_utc.isoformat()
                }
            )
            
    except Exception as e:
        logger.error(f"Error checking daily jobs for {target_partition_key}: {str(e)}")
        # Send error alert to Slack
        error_message = _create_error_alert_message(target_partition_key, str(e), current_utc)
        #_send_slack_alert(error_message, logger)
        
        return dg.SensorResult(
            skip_reason=f"Error checking daily jobs: {str(e)}",
            cursor_result=target_partition_key
        )


def _check_asset_materialization(
    instance: DagsterInstance, 
    asset_key: AssetKey, 
    asset_name: str,
    target_partition_key: str,
    logger
) -> bool:
    """Check if an asset has materialization for the target partition."""
    try:
        # Try multiple approaches to check for materialization
        
        # Method 1: Try fetch_materializations (recommended API)
        try:
            if hasattr(instance, 'fetch_materializations'):
                materializations = instance.fetch_materializations(
                    records_filter=asset_key,
                    limit=50  # Daily assets have fewer partitions than hourly
                )
                
                for materialization in materializations.records:
                    if hasattr(materialization, 'partition_key') and materialization.partition_key == target_partition_key:
                        logger.debug(f"Found materialization for {asset_name} partition {target_partition_key}")
                        return True
                        
        except Exception as e:
            logger.debug(f"Method 1 (fetch_materializations) failed for {asset_name}: {e}")
        
        # Method 2: Try get_event_records with proper filter
        try:
            from dagster import EventRecordsFilter, DagsterEventType
            event_filter = EventRecordsFilter(
                event_type=DagsterEventType.ASSET_MATERIALIZATION,
                asset_key=asset_key
            )
            event_records = instance.get_event_records(event_filter, limit=100)
            
            for event_record in event_records:
                event_log_entry = event_record.event_log_entry
                
                if (hasattr(event_log_entry, 'dagster_event') and 
                    event_log_entry.dagster_event and
                    event_log_entry.dagster_event.is_step_materialization):
                    
                    # Get partition key
                    partition_key = None
                    if hasattr(event_log_entry, 'partition_key'):
                        partition_key = event_log_entry.partition_key
                    elif hasattr(event_log_entry.dagster_event, 'partition_key'):
                        partition_key = event_log_entry.dagster_event.partition_key
                    
                    if partition_key == target_partition_key:
                        logger.debug(f"Found materialization for {asset_name} partition {target_partition_key}")
                        return True
                        
        except Exception as e:
            logger.debug(f"Method 2 (get_event_records) failed for {asset_name}: {e}")
        
        # Method 3: Check asset records
        try:
            asset_record = instance.get_asset_record(asset_key)
            if asset_record and asset_record.asset_entry.last_materialization:
                last_mat = asset_record.asset_entry.last_materialization
                if hasattr(last_mat, 'partition_key') and last_mat.partition_key == target_partition_key:
                    logger.debug(f"Found asset record materialization for {asset_name} partition {target_partition_key}")
                    return True
        except Exception as e:
            logger.debug(f"Method 3 (asset_record) failed for {asset_name}: {e}")
        
        logger.warning(f"No materialization found for {asset_name} partition {target_partition_key}")
        return False
        
    except Exception as e:
        logger.error(f"Error checking materialization for {asset_name}: {str(e)}")
        return False


def _create_consolidated_alert_message(
    missing_assets: List[Dict], 
    partition_key: str, 
    check_time: datetime
) -> str:
    """Create consolidated Slack alert message for all missing partitions."""
    
    # Group by status
    missing_materializations = [a for a in missing_assets if a["status"] == "missing_materialization"]
    check_errors = [a for a in missing_assets if a["status"] == "check_error"]
    
    message = f"""🚨 *Daily Jobs SLA Alert* 🚨

*Partition Date:* `{partition_key}`
*Check Time:* {check_time.strftime('%Y-%m-%d %H:%M:%S')} UTC
*SLA Deadline:* {SLA_DEADLINE_HOUR}:00 UTC

*Summary:* {len(missing_assets)} assets missing partitions"""
    
    if missing_materializations:
        message += f"""

*Missing Materializations ({len(missing_materializations)}):*"""
        for asset in missing_materializations[:10]:  # Limit to first 10 to avoid message length issues
            message += f"\n• `{asset['asset_name']}`"
        
        if len(missing_materializations) > 10:
            message += f"\n• ... and {len(missing_materializations) - 10} more"
    
    if check_errors:
        message += f"""

*Check Errors ({len(check_errors)}):*"""
        for asset in check_errors[:5]:  # Limit errors to first 5
            message += f"\n• `{asset['asset_name']}`: {asset.get('error', 'Unknown error')}"
        
        if len(check_errors) > 5:
            message += f"\n• ... and {len(check_errors) - 5} more errors"
    
    message += f"""

*Action Required:* Check Dagster UI for failed jobs and partitions
*Runbook:* Investigate assets and backfill missing partitions
*Channel:* #{SLACK_CHANNEL}"""
    
    return message


def _create_error_alert_message(partition_key: str, error: str, check_time: datetime) -> str:
    """Create formatted Slack alert message for sensor errors."""
    return f"""⚠️ *Daily Jobs Monitor Error* ⚠️

*Partition:* `{partition_key}`
*Check Time:* {check_time.strftime('%Y-%m-%d %H:%M:%S')} UTC
*Error:* {error}

*Status:* ❌ Daily jobs monitoring failed
*Action Required:* Please investigate monitoring sensor

*Channel:* #{SLACK_CHANNEL}"""


def _send_slack_alert(message: str, logger) -> bool:
    """Send alert message to Slack channel."""
    try:
        slack_url = "https://slack.com/api/chat.postMessage"
        headers = {
            "Authorization": f"Bearer {SLACK_TOKEN}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "channel": f"#{SLACK_CHANNEL}",
            "text": message,
            "username": "Dagster Daily Jobs Monitor",
            "icon_emoji": ":rotating_light:"
        }
        
        response = requests.post(slack_url, headers=headers, json=payload)
        response_data = response.json()
        
        if response_data.get("ok"):
            logger.info(f"Slack alert sent successfully to #{SLACK_CHANNEL}")
            return True
        else:
            logger.error(f"Failed to send Slack alert: {response_data.get('error', 'Unknown error')}")
            return False
            
    except Exception as e:
        logger.error(f"Exception sending Slack alert: {str(e)}")
        return False