#!/usr/bin/env python3
"""
Trigger Airflow DAG for a document already in SeaweedFS.

Usage:
    # Trigger DAG for a file key
    python scripts/trigger_dag.py 1706.03762v7.pdf

    # Trigger with custom document ID
    python scripts/trigger_dag.py path/to/file.xlsx --document-id my-uuid

    # Trigger with custom bucket
    python scripts/trigger_dag.py path/to/file.pdf --bucket my-bucket
"""

import argparse
import os
import sys
import uuid

import requests
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Content type mapping
EXTENSION_TO_CONTENT_TYPE = {
    'xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
    'xls': 'application/vnd.ms-excel',
    'xlsm': 'application/vnd.ms-excel.sheet.macroEnabled.12',
    'docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
    'doc': 'application/msword',
    'pdf': 'application/pdf',
    'pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
    'ppt': 'application/vnd.ms-powerpoint',
}


def get_content_type(file_key: str) -> str:
    """Get content type from file extension."""
    ext = file_key.lower().rsplit('.', 1)[-1] if '.' in file_key else ''
    return EXTENSION_TO_CONTENT_TYPE.get(ext, 'application/octet-stream')


def get_airflow_token(airflow_url: str, username: str, password: str) -> str | None:
    """Get JWT token from Airflow 3.x API."""
    url = f"{airflow_url}/auth/token"
    try:
        response = requests.post(
            url,
            json={'username': username, 'password': password},
            headers={'Content-Type': 'application/json'},
            timeout=30,
        )
        if response.status_code in (200, 201):
            return response.json().get('access_token')
        else:
            print(f"Failed to get token: {response.status_code} - {response.text}")
            return None
    except Exception as e:
        print(f"Failed to get token: {e}")
        return None


def trigger_dag(
    file_key: str,
    bucket: str,
    document_id: str,
    airflow_url: str,
    username: str,
    password: str,
    dag_id: str = 'document_extraction_dag',
) -> bool:
    """Trigger the document_extraction_dag for a file."""
    # Get JWT token
    token = get_airflow_token(airflow_url, username, password)
    if not token:
        return False

    url = f"{airflow_url}/api/v2/dags/{dag_id}/dagRuns"

    content_type = get_content_type(file_key)

    payload = {
        'logical_date': None,
        'conf': {
            'document_id': document_id,
            'bucket': bucket,
            'object_key': file_key,
            'content_type': content_type,
            'size_bytes': 0,  # Unknown, validator will handle
        },
    }

    try:
        response = requests.post(
            url,
            json=payload,
            headers={
                'Content-Type': 'application/json',
                'Authorization': f'Bearer {token}',
            },
            timeout=30,
        )

        if response.status_code in (200, 201):
            data = response.json()
            run_id = data.get('dag_run_id', 'unknown')
            print(f"DAG triggered successfully!")
            print(f"  Run ID: {run_id}")
            print(f"  DAG ID: {dag_id}")
            print(f"  File: s3://{bucket}/{file_key}")
            print(f"  Document ID: {document_id}")
            print(f"  Content Type: {content_type}")
            return True
        else:
            print(f"Failed to trigger DAG: {response.status_code} - {response.text}")
            return False

    except Exception as e:
        print(f"Failed to trigger DAG: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(
        description='Trigger Airflow DAG for a document in SeaweedFS',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    parser.add_argument(
        'file_key',
        help='File key/path in SeaweedFS (e.g., "1706.03762v7.pdf" or "folder/file.xlsx")',
    )
    parser.add_argument(
        '--bucket', '-b',
        default='documents',
        help='S3 bucket name (default: documents)',
    )
    parser.add_argument(
        '--document-id', '-d',
        default=None,
        help='Document ID (default: auto-generated UUID)',
    )
    parser.add_argument(
        '--dag-id',
        default='document_extraction_dag',
        help='DAG ID to trigger (default: document_extraction_dag)',
    )
    parser.add_argument(
        '--airflow-url',
        default=os.getenv('AIRFLOW_URL', 'http://localhost:8081'),
        help='Airflow API URL (default: http://localhost:8081)',
    )
    parser.add_argument(
        '--airflow-user',
        default=os.getenv('_AIRFLOW_WWW_USER_USERNAME', 'admin'),
        help='Airflow username',
    )
    parser.add_argument(
        '--airflow-password',
        default=os.getenv('_AIRFLOW_WWW_USER_PASSWORD', 'admin'),
        help='Airflow password',
    )

    args = parser.parse_args()

    # Generate document ID if not provided
    document_id = args.document_id or str(uuid.uuid4())

    # Trigger DAG
    success = trigger_dag(
        file_key=args.file_key,
        bucket=args.bucket,
        document_id=document_id,
        airflow_url=args.airflow_url,
        username=args.airflow_user,
        password=args.airflow_password,
        dag_id=args.dag_id,
    )

    sys.exit(0 if success else 1)


if __name__ == '__main__':
    main()