#!/usr/bin/env python3
"""
Upload documents to SeaweedFS and trigger Airflow DAG.

This script uploads files or folders to SeaweedFS S3-compatible storage
and optionally triggers the document_extraction_dag for processing.

Usage:
    # Upload a single file
    python scripts/upload_documents.py path/to/file.xlsx

    # Upload a folder
    python scripts/upload_documents.py path/to/folder/

    # Upload and trigger DAG
    python scripts/upload_documents.py path/to/file.xlsx --trigger

    # Upload to specific bucket
    python scripts/upload_documents.py path/to/file.xlsx --bucket my-bucket

    # Dry run (show what would be uploaded)
    python scripts/upload_documents.py path/to/folder/ --dry-run
"""

import argparse
import mimetypes
import os
import sys
import uuid
from pathlib import Path

import boto3
import requests
from botocore.config import Config
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Supported file extensions - PDF, Word, and PowerPoint only
SUPPORTED_EXTENSIONS = {
    '.pdf': 'application/pdf',
    '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
    '.doc': 'application/msword',
    '.docm': 'application/vnd.ms-word.document.macroEnabled.12',
    '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
    '.ppt': 'application/vnd.ms-powerpoint',
    '.xls': 'application/vnd.ms-excel',
    '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
    '.xlsm': 'application/vnd.ms-excel.sheet.macroEnabled.12',
}


def get_s3_client():
    """Create S3 client for SeaweedFS."""
    # Use SEAWEEDFS_ENDPOINT if set, otherwise build from host/port
    # This allows easy override: SEAWEEDFS_ENDPOINT=http://localhost:8333
    endpoint = os.getenv('SEAWEEDFS_ENDPOINT')
    if not endpoint:
        host = os.getenv('SEAWEEDFS_HOST', 'localhost')
        port = os.getenv('SEAWEEDFS_PORT', '8333')
        # Check SSL settings - either env var enables HTTPS
        use_ssl = os.getenv('SEAWEEDFS_USE_SSL', '').lower() == 'true' or \
                  os.getenv('S3_USE_HTTPS', '').lower() == 'true'
        protocol = 'https' if use_ssl else 'http'
        endpoint = f"{protocol}://{host}:{port}"

    access_key = os.getenv('SEAWEEDFS_ACCESS_KEY_ID', 'dummy')
    secret_key = os.getenv('SEAWEEDFS_SECRET_ACCESS_KEY', 'dummy')

    return boto3.client(
        's3',
        endpoint_url=endpoint,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
        config=Config(signature_version='s3v4'),
        region_name='us-east-1',
    )


def ensure_bucket_exists(s3_client, bucket_name: str) -> bool:
    """Ensure the bucket exists, create if not."""
    try:
        s3_client.head_bucket(Bucket=bucket_name)
        return True
    except s3_client.exceptions.ClientError as e:
        error_code = e.response.get('Error', {}).get('Code', '')
        if error_code == '404':
            try:
                s3_client.create_bucket(Bucket=bucket_name)
                print(f"Created bucket: {bucket_name}")
                return True
            except Exception as create_err:
                print(f"Failed to create bucket: {create_err}")
                return False
        else:
            print(f"Error checking bucket: {e}")
            return False


def get_content_type(file_path: Path) -> str:
    """Get content type for a file."""
    ext = file_path.suffix.lower()
    if ext in SUPPORTED_EXTENSIONS:
        return SUPPORTED_EXTENSIONS[ext]

    # Fallback to mimetypes
    content_type, _ = mimetypes.guess_type(str(file_path))
    return content_type or 'application/octet-stream'


def is_supported_file(file_path: Path) -> bool:
    """Check if file is supported for extraction."""
    return file_path.suffix.lower() in SUPPORTED_EXTENSIONS


def upload_file(
    s3_client,
    file_path: Path,
    bucket: str,
    key_prefix: str = '',
    dry_run: bool = False,
    tenant_id: str = '',
    node_id: str = '',
) -> dict | None:
    """
    Upload a single file to SeaweedFS.

    Returns dict with upload info or None if skipped/failed.
    """
    if not file_path.is_file():
        print(f"  Skipping (not a file): {file_path}")
        return None

    if not is_supported_file(file_path):
        print(f"  Skipping (unsupported): {file_path}")
        return None

    # Generate document ID
    document_id = str(uuid.uuid4())

    # Generate tenant-scoped object key: {tenant_id}/{node_id}/{document_id}/{filename}
    if tenant_id and node_id:
        base_key = f"{tenant_id}/{node_id}/{document_id}/{file_path.name}"
    elif key_prefix:
        base_key = f"{key_prefix}/{file_path.name}"
    else:
        base_key = file_path.name
    object_key = base_key.lstrip('/')

    # Get content type
    content_type = get_content_type(file_path)

    if dry_run:
        print(f"  [DRY RUN] Would upload: {file_path.name}")
        print(f"            -> s3://{bucket}/{object_key}")
        print(f"            Content-Type: {content_type}")
        print(f"            Document ID: {document_id}")
        print(f"            Tenant ID: {tenant_id}")
        print(f"            Node ID: {node_id}")
        return {
            'document_id': document_id,
            'bucket': bucket,
            'object_key': object_key,
            'content_type': content_type,
            'file_path': str(file_path),
            'size_bytes': file_path.stat().st_size,
            'tenant_id': tenant_id,
            'node_id': node_id,
        }

    try:
        # Upload file
        with open(file_path, 'rb') as f:
            s3_client.put_object(
                Bucket=bucket,
                Key=object_key,
                Body=f,
                ContentType=content_type,
                Metadata={'document_id': document_id},
            )

        file_size = file_path.stat().st_size
        print(f"  Uploaded: {file_path.name} ({file_size:,} bytes)")
        print(f"         -> s3://{bucket}/{object_key}")

        return {
            'document_id': document_id,
            'bucket': bucket,
            'object_key': object_key,
            'content_type': content_type,
            'file_path': str(file_path),
            'size_bytes': file_size,
            'tenant_id': tenant_id,
            'node_id': node_id,
        }

    except Exception as e:
        print(f"  Failed to upload {file_path.name}: {e}")
        return None


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(upload_info: dict, airflow_url: str, username: str, password: str) -> bool:
    """Trigger the document_extraction_dag for a uploaded file."""
    # Get JWT token for Airflow 3.x API
    token = get_airflow_token(airflow_url, username, password)
    if not token:
        return False

    dag_id = 'document_extraction_dag'
    # Airflow 3.x uses /api/v2/ endpoint
    url = f"{airflow_url}/api/v2/dags/{dag_id}/dagRuns"

    payload = {
        'logical_date': None,  # Airflow 3.x: use None for event-triggered runs
        'conf': {
            'document_id': upload_info['document_id'],
            'bucket': upload_info['bucket'],
            'object_key': upload_info['object_key'],
            'tenant_id': upload_info.get('tenant_id', ''),
            'node_id': upload_info.get('node_id', ''),
        }
    }

    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):
            run_id = response.json().get('dag_run_id', 'unknown')
            print(f"  Triggered DAG run: {run_id}")
            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 collect_files(path: Path, recursive: bool = True) -> list[Path]:
    """Collect all supported files from path."""
    if path.is_file():
        if is_supported_file(path):
            return [path]
        return []

    if path.is_dir():
        if recursive:
            files = []
            for ext in SUPPORTED_EXTENSIONS:
                files.extend(path.rglob(f'*{ext}'))
            return sorted(files)
        else:
            files = [
                f for f in path.iterdir()
                if f.is_file() and is_supported_file(f)
            ]
            return sorted(files)

    return []


def main():
    parser = argparse.ArgumentParser(
        description='Upload documents to SeaweedFS and trigger Airflow DAG',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    parser.add_argument(
        'path',
        type=Path,
        help='File or folder to upload',
    )
    parser.add_argument(
        '--bucket', '-b',
        default='documents',
        help='S3 bucket name (default: documents)',
    )
    parser.add_argument(
        '--prefix', '-p',
        default='',
        help='Key prefix for uploaded files',
    )
    parser.add_argument(
        '--trigger', '-t',
        action='store_true',
        help='Trigger document_extraction_dag after upload',
    )
    parser.add_argument(
        '--dry-run', '-n',
        action='store_true',
        help='Show what would be uploaded without actually uploading',
    )
    parser.add_argument(
        '--no-recursive', '-R',
        action='store_true',
        help='Do not recurse into subdirectories',
    )
    parser.add_argument(
        '--tenant-id',
        default='00000000-0000-0000-0000-000000000001',
        help='Tenant UUID for data isolation (default: pilot tenant)',
    )
    parser.add_argument(
        '--node-id',
        default='00000000-0000-0000-0000-000000000001',
        help='Node UUID for data isolation (default: pilot node)',
    )
    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()

    # Validate path
    if not args.path.exists():
        print(f"Error: Path does not exist: {args.path}")
        sys.exit(1)

    # Collect files
    files = collect_files(args.path, recursive=not args.no_recursive)

    if not files:
        print(f"No supported files found in: {args.path}")
        print(f"Supported extensions: {', '.join(SUPPORTED_EXTENSIONS.keys())}")
        sys.exit(1)

    print(f"Found {len(files)} file(s) to upload")
    print(f"Bucket: {args.bucket}")
    print(f"Tenant ID: {args.tenant_id}")
    print(f"Node ID: {args.node_id}")
    if args.prefix:
        print(f"Prefix: {args.prefix}")
    print()

    # Create S3 client
    if not args.dry_run:
        s3_client = get_s3_client()

        # Ensure bucket exists
        if not ensure_bucket_exists(s3_client, args.bucket):
            print("Failed to ensure bucket exists")
            sys.exit(1)
    else:
        s3_client = None

    # Upload files
    uploaded = []
    for file_path in files:
        result = upload_file(
            s3_client,
            file_path,
            args.bucket,
            args.prefix,
            args.dry_run,
            tenant_id=args.tenant_id,
            node_id=args.node_id,
        )
        if result:
            uploaded.append(result)

    print()
    print(f"Uploaded: {len(uploaded)}/{len(files)} files")

    # Trigger DAGs if requested
    if args.trigger and uploaded and not args.dry_run:
        print()
        print("Triggering DAG runs...")
        triggered = 0
        for upload_info in uploaded:
            if trigger_dag(
                upload_info,
                args.airflow_url,
                args.airflow_user,
                args.airflow_password,
            ):
                triggered += 1
        print(f"Triggered: {triggered}/{len(uploaded)} DAG runs")
    elif args.trigger and args.dry_run:
        print()
        print(f"[DRY RUN] Would trigger {len(uploaded)} DAG run(s)")


if __name__ == '__main__':
    main()