#!/usr/bin/env python3
"""
Script to update DBT Cloud notification configurations for Slack channels.

This script updates notification settings to include all jobs from specified projects
in the on_failure notification list for designated Slack channels.

Usage:
    python update_dbt_notifications.py --dry-run  # Preview changes without applying
    python update_dbt_notifications.py            # Apply changes
"""

import argparse
import os
import sys
from typing import Dict, List, Optional
import requests


# Configuration: Map projects to Slack channels and job exclusions
PROJECT_CONFIG = {
    "Analytics Orchard": {
        "slack_channel": "#dbt_jobs_monitoring",
        "excluded_jobs": ["New Slim CI Job"],
    },
    "Analytics Orchard Strategy": {
        "slack_channel": "#strategy-team-dbt-alerts",
        "excluded_jobs": [],
    },
}


class DBTCloudAPI:
    """Client for interacting with DBT Cloud API."""

    def __init__(self, api_key: str, account_id: str, base_url: str = "https://cloud.getdbt.com/api/v2"):
        self.api_key = api_key
        self.account_id = account_id
        self.base_url = base_url
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Token {api_key}",
            "Content-Type": "application/json",
        })

    def _get(self, endpoint: str, params: Optional[Dict] = None) -> Dict:
        """Make a GET request to the DBT Cloud API."""
        url = f"{self.base_url}/{endpoint}"
        response = self.session.get(url, params=params)
        response.raise_for_status()
        return response.json()

    def _post(self, endpoint: str, data: Dict) -> Dict:
        """Make a POST request to the DBT Cloud API."""
        url = f"{self.base_url}/{endpoint}"
        response = self.session.post(url, json=data)
        response.raise_for_status()
        return response.json()

    def _patch(self, endpoint: str, data: Dict) -> Dict:
        """Make a PATCH request to the DBT Cloud API."""
        url = f"{self.base_url}/{endpoint}"
        response = self.session.patch(url, json=data)
        response.raise_for_status()
        return response.json()

    def get_projects(self) -> List[Dict]:
        """Get all projects in the account."""
        response = self._get(f"accounts/{self.account_id}/projects/")
        return response.get("data", [])

    def get_project_by_name(self, project_name: str) -> Optional[Dict]:
        """Get a project by its name."""
        projects = self.get_projects()
        for project in projects:
            if project.get("name") == project_name:
                return project
        return None

    def get_jobs(self, project_id: int) -> List[Dict]:
        """Get all jobs for a project, handling pagination.

        Args:
            project_id: The project ID to fetch jobs for

        Returns:
            List of all job objects for the project
        """
        all_jobs = []
        limit = 100  # Maximum items per page
        offset = 0

        while True:
            params = {
                "project_id": project_id,
                "limit": limit,
                "offset": offset
            }
            response = self._get(f"accounts/{self.account_id}/jobs/", params=params)

            jobs = response.get("data", [])
            all_jobs.extend(jobs)

            # Check if there are more pages
            # The API returns extra.pagination info or we can check if we got fewer items than limit
            extra = response.get("extra", {})
            pagination = extra.get("pagination", {})
            total_count = pagination.get("total_count")

            # If we know the total count, check if we've fetched everything
            if total_count is not None:
                if len(all_jobs) >= total_count:
                    break
            else:
                # If no total_count, check if we got fewer items than requested
                if len(jobs) < limit:
                    break

            offset += limit

        return all_jobs

    def get_notifications(self) -> List[Dict]:
        """Get all active notifications for the account.

        Returns:
            List of notification objects
        """
        params = {"state": "active"}  # Only get active notifications
        response = self._get(f"accounts/{self.account_id}/notifications/", params=params)
        return response.get("data", [])

    def get_notification_by_channel(self, channel_name: str) -> Optional[Dict]:
        """Find notification by Slack channel name.

        Args:
            channel_name: Slack channel name (e.g., "#dbt_jobs_monitoring")

        Returns:
            Notification object if found, None otherwise
        """
        notifications = self.get_notifications()

        # Filter by channel name and return the first match
        for notification in notifications:
            if notification.get("slack_channel_name") == channel_name:
                return notification

        return None

    def update_notification(self, notification_id: int, notification_data: Dict) -> Dict:
        """Update a notification.

        Args:
            notification_id: The notification ID
            notification_data: Full notification object with updated fields

        Returns:
            Updated notification object
        """
        return self._post(
            f"accounts/{self.account_id}/notifications/{notification_id}/",
            data=notification_data
        )


def update_notifications(api: DBTCloudAPI, dry_run: bool = False) -> None:
    """
    Update notification configurations for all projects in the mapping.

    Args:
        api: DBTCloudAPI client instance
        dry_run: If True, only print what would be done without making changes
    """
    print("=" * 80)
    print("DBT Cloud Notification Configuration Update")
    print("=" * 80)
    print(f"Mode: {'DRY RUN (no changes will be made)' if dry_run else 'LIVE (changes will be applied)'}")
    print("=" * 80)
    print()

    for project_name, config in PROJECT_CONFIG.items():
        slack_channel = config["slack_channel"]
        excluded_jobs = config.get("excluded_jobs", [])

        print(f"\n{'─' * 80}")
        print(f"Processing: {project_name} → {slack_channel}")
        print(f"{'─' * 80}")

        # Get the project
        project = api.get_project_by_name(project_name)
        if not project:
            print(f"❌ ERROR: Project '{project_name}' not found")
            continue

        project_id = project["id"]
        print(f"✓ Found project ID: {project_id}")

        # Get all jobs for the project
        jobs = api.get_jobs(project_id)
        if not jobs:
            print(f"⚠️  WARNING: No jobs found in project '{project_name}'")
            continue

        # Filter out excluded jobs
        filtered_jobs = [job for job in jobs if job["name"] not in excluded_jobs]

        if excluded_jobs:
            excluded_count = len(jobs) - len(filtered_jobs)
            if excluded_count > 0:
                print(f"ℹ️  Excluding {excluded_count} job(s): {', '.join(excluded_jobs)}")

        # Create mapping of job ID to job name
        job_id_to_name = {job["id"]: job["name"] for job in filtered_jobs}
        job_ids = [job["id"] for job in filtered_jobs]

        # Helper function to format job IDs as names for display
        def format_jobs(job_id_list):
            return [f"{job_id_to_name.get(jid, f'Unknown Job {jid}')} (ID: {jid})" for jid in job_id_list]

        print(f"✓ Found {len(jobs)} jobs:")

        # Get notification configuration for the Slack channel
        notification = api.get_notification_by_channel(slack_channel)

        if notification:
            notification_id = notification["id"]
            print(f"✓ Found existing notification ID: {notification_id}")

            # Get current on_failure list
            current_on_failure = notification.get("on_failure", [])

            # Merge with new job IDs (avoid duplicates)
            updated_on_failure = sorted(list(set(current_on_failure + job_ids)))
            new_jobs = [job_id for job_id in job_ids if job_id not in current_on_failure]

            # Check if update is needed
            if set(current_on_failure) == set(updated_on_failure):
                print(f"ℹ️  No changes needed - all jobs already in on_failure list")
                continue

            print(f"  New jobs to add ({len(new_jobs)}):")
            for job_display in format_jobs(new_jobs):
                print(f"    - {job_display}")

            if dry_run:
                print(f"🔍 DRY RUN: Would update notification {notification_id}")
                print(f"   Would add {len(new_jobs)} new job(s) to on_failure list")
            else:
                # Update the notification with the modified on_failure list
                # The API requires the full notification object
                notification["on_failure"] = updated_on_failure
                api.update_notification(notification_id, notification)
                print(f"✅ Updated notification {notification_id}")
                print(f"   Added {len(new_jobs)} new job(s) to on_failure list")
        else:
            print(f"⚠️  WARNING: No existing notification found for {slack_channel}")
            print(f"   You need to create the Slack notification manually in DBT Cloud UI first.")
            print(f"   To create:")
            print(f"   1. Go to Account Settings > Notifications")
            print(f"   2. Click 'Add Notification'")
            print(f"   3. Select Slack and choose channel: {slack_channel}")
            print(f"   4. Save the notification")
            print(f"   5. Re-run this script")

    print("\n" + "=" * 80)
    if dry_run:
        print("DRY RUN COMPLETE - No changes were made")
    else:
        print("UPDATE COMPLETE")
    print("=" * 80)


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Update DBT Cloud notification configurations",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Environment Variables Required:
  DBT_CLOUD_API_KEY      - Your DBT Cloud API key
  DBT_CLOUD_ACCOUNT_ID   - Your DBT Cloud account ID (defaults to 21748)

Examples:
  # Dry run to see what would change
  python update_dbt_notifications.py --dry-run

  # Apply the changes
  python update_dbt_notifications.py
        """
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Preview changes without applying them"
    )
    parser.add_argument(
        "--api-key",
        help="DBT Cloud API key (or set DBT_CLOUD_API_KEY env var)"
    )
    parser.add_argument(
        "--account-id",
        help="DBT Cloud account ID (or set DBT_CLOUD_ACCOUNT_ID env var, defaults to 21748)"
    )

    args = parser.parse_args()

    # Get API credentials
    api_key = args.api_key or os.environ.get("DBT_CLOUD_API_KEY")
    account_id = args.account_id or os.environ.get("DBT_CLOUD_ACCOUNT_ID") or "21748"

    if not api_key:
        print("❌ ERROR: DBT Cloud API key is required", file=sys.stderr)
        print("   Set DBT_CLOUD_API_KEY environment variable or use --api-key", file=sys.stderr)
        sys.exit(1)

    try:
        # Initialize API client
        api = DBTCloudAPI(api_key, account_id)

        # Run the update
        update_notifications(api, dry_run=args.dry_run)

    except requests.exceptions.HTTPError as e:
        print(f"\n❌ API Error: {e}", file=sys.stderr)
        print(f"   Response: {e.response.text}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()