#!/usr/bin/env -S uv --quiet run --script
# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "requests",
# ]
# ///

import sys
import requests
import argparse
import logging
import asyncio
import re
import json
from asyncio import Semaphore
from typing import List, Dict, Optional, Tuple

logging.basicConfig(
    level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def parse_curl_command(curl_command: str) -> Tuple[Dict[str, str], Optional[str]]:
    """
    Parse a curl command to extract headers, authentication token, and cookies.

    Args:
        curl_command: Full curl command string

    Returns:
        Tuple of (headers_dict, auth_token, cookies_string)
    """
    headers = {}
    auth_token = None
    cookies = None

    # Extract all -H headers using multiple patterns to handle different quote styles
    # Pattern 1: -H 'header: value' (single quotes)
    # Pattern 2: -H "header: value" (double quotes)
    # We need to match the entire value including spaces

    # Try to match -H with single quotes first
    single_quote_pattern = r"-H\s+'([^']+)'"
    for match in re.finditer(single_quote_pattern, curl_command):
        header_line = match.group(1)
        if ':' in header_line:
            header_name, header_value = header_line.split(':', 1)
            headers[header_name.strip()] = header_value.strip()

    # Then try double quotes (may override if both exist)
    double_quote_pattern = r'-H\s+"([^"]+)"'
    for match in re.finditer(double_quote_pattern, curl_command):
        header_line = match.group(1)
        if ':' in header_line:
            header_name, header_value = header_line.split(':', 1)
            headers[header_name.strip()] = header_value.strip()

    # Extract authentication token from --data-raw payload
    # Try both single and double quotes
    data_raw_patterns = [
        r"--data-raw\s+'({.*?})'",
        r'--data-raw\s+"({.*?})"',
        r'--data-raw\s+\'({[^\']+})\'',
        r'--data-raw\s+"({[^"]+})"',
    ]

    for pattern in data_raw_patterns:
        data_match = re.search(pattern, curl_command, re.DOTALL)
        if data_match:
            try:
                json_str = data_match.group(1)
                # Unescape common escape sequences
                json_str = json_str.replace('\\"', '"').replace("\\'", "'")
                payload = json.loads(json_str)
                auth_token = payload.get('_authentication_token')
                if auth_token:
                    break
            except json.JSONDecodeError as e:
                logger.debug(f'Failed to parse JSON with pattern {pattern}: {e}')
                continue

    # Extract cookies from --cookie or -b flag
    cookie_patterns = [
        r"--cookie\s+'([^']+)'",
        r'--cookie\s+"([^"]+)"',
        r"-b\s+'([^']+)'",
        r'-b\s+"([^"]+)"',
    ]

    for pattern in cookie_patterns:
        cookie_match = re.search(pattern, curl_command)
        if cookie_match:
            cookies = cookie_match.group(1)
            break

    # Ensure Content-Type is set
    if 'Content-Type' not in headers:
        headers['Content-Type'] = 'application/json'

    logger.info(f'Parsed {len(headers)} headers from curl command')
    logger.debug('Parsed headers:')
    for key, value in headers.items():
        # Mask sensitive values in logs
        if 'token' in key.lower() or 'auth' in key.lower():
            logger.debug(f'  {key}: ***REDACTED***')
        else:
            logger.debug(
                f'  {key}: {value[:50]}...' if len(value) > 50 else f'  {key}: {value}'
            )

    if auth_token:
        logger.info('Successfully extracted authentication token from curl command')
        logger.debug(f'Token: {auth_token[:10]}...{auth_token[-10:]}')
    else:
        logger.warning('Could not extract authentication token from curl command')

    if cookies:
        logger.info('Successfully extracted cookies from curl command')
        logger.debug(
            f'Cookies: {cookies[:50]}...'
            if len(cookies) > 50
            else f'Cookies: {cookies}'
        )
    else:
        logger.warning(
            'No cookies found in curl command - this may cause authorization failures'
        )

    return headers, auth_token, cookies


class DatadogScanningDisabler:
    """Disables Datadog code scanning for archived GitHub repositories."""

    # The 4 different scanning types to disable
    SCAN_TYPES = [
        {'attribute': 'sa_enabled', 'name': 'Static Analysis'},
        {'attribute': 'sca_enabled', 'name': 'Software Composition Analysis'},
        {'attribute': 'secrets_enabled', 'name': 'Secrets Scanning'},
        {'attribute': 'iac_scan_enabled', 'name': 'IaC Scanning'},
    ]

    def __init__(
        self,
        github_token: str,
        org_name: str,
        datadog_token: Optional[str] = None,
        datadog_headers: Optional[Dict[str, str]] = None,
        datadog_cookies: Optional[str] = None,
        max_concurrent: int = 5,
        delay_seconds: float = 0.5,
        dry_run: bool = False,
    ):
        """
        Initialize the Datadog scanning disabler.

        Args:
            github_token: GitHub personal access token
            org_name: Name of the GitHub organization
            datadog_token: Datadog CSRF token (x-csrf-token from browser session)
            datadog_headers: Custom headers to use for Datadog API requests (from curl)
            datadog_cookies: Cookie string from browser session (critical for auth)
            max_concurrent: Maximum number of concurrent Datadog API calls
            delay_seconds: Delay between Datadog API calls
            dry_run: If True, only preview actions without making changes
        """
        self.github_token = github_token
        self.datadog_token = datadog_token
        self.datadog_cookies = datadog_cookies
        self.org_name = org_name
        self.max_concurrent = max_concurrent
        self.delay_seconds = delay_seconds
        self.dry_run = dry_run

        self.github_headers = {
            'Authorization': f'token {github_token}',
            'Accept': 'application/vnd.github.v3+json',
        }
        self.github_api_base = 'https://api.github.com'
        self.datadog_api_url = 'https://sonymusic-pde.datadoghq.com/api/v2/source-code/repositories/cwp_settings'

        # Use custom headers if provided, otherwise create default headers
        if datadog_headers:
            self.datadog_headers = datadog_headers
        elif datadog_token:
            self.datadog_headers = {
                'Content-Type': 'application/json',
                'x-csrf-token': datadog_token,
            }
        else:
            raise ValueError('Either datadog_token or datadog_headers must be provided')

        self.semaphore = Semaphore(max_concurrent)
        self.errors: List[str] = []
        self.success_count = 0
        self.failed_count = 0

    def get_archived_repositories(self) -> List[Dict]:
        """Fetch all archived repositories for the organization."""
        repos = []
        page = 1

        logger.info(f'Fetching archived repositories from {self.org_name}...')

        while True:
            url = f'{self.github_api_base}/orgs/{self.org_name}/repos'
            params = {'page': page, 'per_page': 100}

            try:
                response = requests.get(
                    url, headers=self.github_headers, params=params, timeout=30
                )
                response.raise_for_status()

                page_repos = response.json()
                if not page_repos:
                    break

                # Filter for archived repositories only
                archived_repos = [
                    repo for repo in page_repos if repo.get('archived', False)
                ]
                repos.extend(archived_repos)

                logger.info(
                    f'Page {page}: Found {len(archived_repos)} archived repos (out of {len(page_repos)} total)'
                )
                page += 1

            except requests.exceptions.RequestException as e:
                logger.error(f'Failed to fetch repositories from GitHub: {e}')
                raise

        logger.info(f'Total archived repositories found: {len(repos)}')
        return repos

    async def disable_scanning_for_repo(self, repo: Dict) -> bool:
        """
        Disable all Datadog scanning types for a single repository.

        Args:
            repo: Repository information from GitHub API

        Returns:
            True if all scanning types were disabled successfully, False otherwise
        """
        async with self.semaphore:
            repo_name = repo['name']
            repo_full_name = f'{self.org_name}/{repo_name}'
            repo_id = f'github.com/{repo_full_name}'

            logger.info(f'Processing repository: {repo_full_name}')

            if self.dry_run:
                logger.info(
                    f'[DRY RUN] Would disable {len(self.SCAN_TYPES)} scanning types for {repo_full_name}'
                )
                return True

            repo_success = True
            first_request = True

            for scan_type in self.SCAN_TYPES:
                try:
                    # Prepare the request payload
                    payload = {
                        'data': {
                            'type': 'source_code_repository_cwp_settings',
                            'attributes': {scan_type['attribute']: False},
                            'id': repo_id,
                        },
                        '_authentication_token': self.datadog_token,
                    }

                    # Prepare headers for this request (add cookies if available)
                    request_headers = self.datadog_headers.copy()
                    if self.datadog_cookies:
                        request_headers['Cookie'] = self.datadog_cookies

                    # Log details for first request to help with debugging
                    if first_request:
                        logger.debug(f'Making PATCH request to: {self.datadog_api_url}')
                        logger.debug(
                            f'Request payload: {json.dumps(payload, indent=2)}'
                        )
                        logger.debug('Request headers:')
                        for key, value in request_headers.items():
                            if (
                                'token' in key.lower()
                                or 'auth' in key.lower()
                                or 'cookie' in key.lower()
                            ):
                                logger.debug(f'  {key}: ***REDACTED***')
                            else:
                                logger.debug(
                                    f'  {key}: {value[:100]}...'
                                    if len(value) > 100
                                    else f'  {key}: {value}'
                                )
                        if self.datadog_cookies:
                            logger.debug('  Cookies: ***INCLUDED (redacted)***')
                        first_request = False

                    # Make the PATCH request with stored headers and cookies
                    response = requests.patch(
                        self.datadog_api_url,
                        json=payload,
                        headers=request_headers,
                        timeout=30,
                    )

                    if response.status_code in [200, 201, 204]:
                        logger.info(
                            f'  ✓ Disabled {scan_type["name"]} for {repo_full_name}'
                        )
                    else:
                        error_msg = f'  ✗ Failed to disable {scan_type["name"]} for {repo_full_name}: HTTP {response.status_code}'
                        logger.error(error_msg)
                        logger.debug(f'Response headers: {dict(response.headers)}')
                        logger.debug(f'Response body: {response.text[:500]}')
                        self.errors.append(f'{error_msg} - {response.text[:200]}')
                        repo_success = False

                    # Throttle requests to avoid overwhelming the API
                    await asyncio.sleep(self.delay_seconds)

                except requests.exceptions.RequestException as e:
                    error_msg = f'  ✗ Request failed for {scan_type["name"]} on {repo_full_name}: {str(e)}'
                    logger.error(error_msg)
                    self.errors.append(error_msg)
                    repo_success = False

                except Exception as e:
                    error_msg = f'  ✗ Unexpected error for {scan_type["name"]} on {repo_full_name}: {str(e)}'
                    logger.error(error_msg)
                    self.errors.append(error_msg)
                    repo_success = False

            if repo_success:
                self.success_count += 1
            else:
                self.failed_count += 1

            return repo_success

    async def process_all_repositories(self):
        """Process all archived repositories and disable Datadog scanning."""
        try:
            # Fetch archived repositories from GitHub
            repos = self.get_archived_repositories()

            if not repos:
                logger.warning('No archived repositories found.')
                return

            if self.dry_run:
                logger.info(
                    f'\n[DRY RUN MODE] Would process {len(repos)} archived repositories'
                )
                logger.info('Repositories that would be affected:')
                for repo in repos:
                    logger.info(f'  - {repo["name"]}')
                logger.info(
                    f'\nEach repository would have {len(self.SCAN_TYPES)} scanning types disabled.'
                )
                return

            # Process all repositories concurrently with rate limiting
            logger.info(
                f'\nStarting to disable scanning for {len(repos)} archived repositories...'
            )
            logger.info(
                f'Concurrency limit: {self.max_concurrent}, Delay between calls: {self.delay_seconds}s\n'
            )

            tasks = [self.disable_scanning_for_repo(repo) for repo in repos]
            await asyncio.gather(*tasks)

            # Print summary
            logger.info('\n' + '=' * 60)
            logger.info('SUMMARY')
            logger.info('=' * 60)
            logger.info(f'Total repositories processed: {len(repos)}')
            logger.info(f'Successfully processed: {self.success_count}')
            logger.info(f'Failed: {self.failed_count}')

            if self.errors:
                logger.error(f'\nTotal errors: {len(self.errors)}')
                logger.error('\nError details:')
                for error in self.errors:
                    logger.error(error)
                sys.exit(1)
            else:
                logger.info('\n✓ All repositories processed successfully!')

        except Exception as e:
            logger.error(f'Failed to process repositories: {e}')
            sys.exit(1)


async def async_main():
    parser = argparse.ArgumentParser(
        description='Disable Datadog code scanning for archived GitHub repositories',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Example usage:
  # Using curl command (recommended):
  %(prog)s --github-token ghp_xxx --org theorchard --from-curl 'curl "https://..." -H "baggage: ..." ...'
  
  # Using manual token:
  %(prog)s --github-token ghp_xxx --datadog-token xxx --org theorchard
  
  # Dry run to preview:
  %(prog)s --github-token ghp_xxx --org theorchard --from-curl 'curl ...' --dry-run
        """,
    )

    parser.add_argument(
        '--github-token', required=True, help='GitHub personal access token'
    )
    parser.add_argument('--org', required=True, help='GitHub organization name')
    parser.add_argument(
        '--from-curl',
        help='Full curl command from browser Network tab (extracts headers and token automatically)',
    )
    parser.add_argument(
        '--datadog-token',
        help='Datadog CSRF token (alternative to --from-curl)',
    )
    parser.add_argument(
        '--concurrent',
        type=int,
        default=5,
        help='Maximum number of concurrent Datadog API calls (default: 5)',
    )
    parser.add_argument(
        '--delay',
        type=float,
        default=0.5,
        help='Delay in seconds between Datadog API calls (default: 0.5)',
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Preview which repositories would be affected without making changes',
    )
    parser.add_argument(
        '--verbose',
        action='store_true',
        help='Enable verbose debug logging',
    )

    args = parser.parse_args()

    # Set logging level based on verbose flag
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
        logger.debug('Verbose logging enabled')

    # Validate that either --from-curl or --datadog-token is provided
    if not args.from_curl and not args.datadog_token:
        parser.error('Either --from-curl or --datadog-token must be provided')

    if args.from_curl and args.datadog_token:
        parser.error('Cannot use both --from-curl and --datadog-token, choose one')

    # Parse curl command if provided
    datadog_headers = None
    datadog_token = args.datadog_token
    datadog_cookies = None

    if args.from_curl:
        logger.info('Parsing curl command to extract headers, token, and cookies...')
        datadog_headers, extracted_token, datadog_cookies = parse_curl_command(
            args.from_curl
        )

        if not extracted_token:
            logger.error('Failed to extract authentication token from curl command')
            sys.exit(1)

        datadog_token = extracted_token
        logger.info('Successfully parsed curl command\n')

    if args.dry_run:
        logger.info('Running in DRY RUN mode - no changes will be made\n')

    disabler = DatadogScanningDisabler(
        github_token=args.github_token,
        org_name=args.org,
        datadog_token=datadog_token,
        datadog_headers=datadog_headers,
        datadog_cookies=datadog_cookies,
        max_concurrent=args.concurrent,
        delay_seconds=args.delay,
        dry_run=args.dry_run,
    )

    await disabler.process_all_repositories()


def main():
    asyncio.run(async_main())


if __name__ == '__main__':
    main()