#!/usr/bin/env python3
"""
Batch upload documents to SeaweedFS and trigger Airflow DAG with progress monitoring.

Uploads all supported files from a folder to SeaweedFS S3-compatible storage,
triggers the document_extraction_dag for each file, and monitors DAG run status.
Progress is tracked in a JSON file that updates in real-time.

Usage:
    # Upload and trigger DAG for all files in a folder
    python scripts/batch_upload_and_monitor.py path/to/folder/

    # Upload only (no DAG trigger)
    python scripts/batch_upload_and_monitor.py path/to/folder/ --no-trigger

    # Resume monitoring from a previous progress file
    python scripts/batch_upload_and_monitor.py --resume progress_20240101_120000.json

    # Custom bucket and prefix
    python scripts/batch_upload_and_monitor.py path/to/folder/ --bucket my-bucket --prefix inbox/

    # Dry run
    python scripts/batch_upload_and_monitor.py path/to/folder/ --dry-run

    # Set polling interval for DAG monitoring
    python scripts/batch_upload_and_monitor.py path/to/folder/ --poll-interval 30
"""

import argparse
import json
import os
import sys
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

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

load_dotenv()

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',
}

# DAG run terminal states
TERMINAL_STATES = {'success', 'failed'}
# States that mean the run is still going
RUNNING_STATES = {'queued', 'running'}


class ProgressTracker:
    """Track upload and DAG run progress in a JSON file."""

    def __init__(self, progress_file: str):
        self.progress_file = progress_file
        self.data = {
            'created_at': datetime.now(timezone.utc).isoformat(),
            'updated_at': None,
            'summary': {
                'total': 0,
                'pending': 0,
                'uploading': 0,
                'uploaded': 0,
                'dag_triggered': 0,
                'running': 0,
                'success': 0,
                'failed': 0,
                'skipped': 0,
            },
            'files': {},
        }

    def load(self) -> bool:
        """Load progress from file. Returns True if loaded successfully."""
        try:
            with open(self.progress_file) as f:
                self.data = json.load(f)
            return True
        except (FileNotFoundError, json.JSONDecodeError):
            return False

    def save(self):
        """Save progress to file."""
        self.data['updated_at'] = datetime.now(timezone.utc).isoformat()
        self._recalculate_summary()
        with open(self.progress_file, 'w') as f:
            json.dump(self.data, f, indent=2, default=str)

    def _recalculate_summary(self):
        """Recalculate summary counts from file entries."""
        summary = {
            'total': 0,
            'pending': 0,
            'uploading': 0,
            'uploaded': 0,
            'dag_triggered': 0,
            'running': 0,
            'success': 0,
            'failed': 0,
            'skipped': 0,
        }
        for entry in self.data['files'].values():
            summary['total'] += 1
            status = entry.get('status', 'pending')
            if status in summary:
                summary[status] += 1
        self.data['summary'] = summary

    def add_file(self, file_path: str, file_name: str, size_bytes: int):
        """Add a file entry with pending status."""
        self.data['files'][file_name] = {
            'file_path': file_path,
            'file_name': file_name,
            'size_bytes': size_bytes,
            'status': 'pending',
            'document_id': None,
            'bucket': None,
            'object_key': None,
            'dag_run_id': None,
            'error': None,
            'upload_time': None,
            'trigger_time': None,
            'completion_time': None,
        }

    def update_status(self, file_name: str, status: str, **kwargs):
        """Update the status and extra fields for a file."""
        if file_name in self.data['files']:
            self.data['files'][file_name]['status'] = status
            for key, value in kwargs.items():
                self.data['files'][file_name][key] = value
            self.save()

    def get_files_by_status(self, *statuses: str) -> list[dict]:
        """Get all file entries matching any of the given statuses."""
        return [
            entry for entry in self.data['files'].values()
            if entry.get('status') in statuses
        ]

    def print_summary(self):
        """Print a formatted summary table."""
        s = self.data['summary']
        print('\n' + '=' * 60)
        print('  PROGRESS SUMMARY')
        print('=' * 60)
        print(f"  Total files:    {s['total']}")
        print(f"  Pending:        {s['pending']}")
        print(f"  Uploading:      {s['uploading']}")
        print(f"  Uploaded:       {s['uploaded']}")
        print(f"  DAG triggered:  {s['dag_triggered']}")
        print(f"  Running:        {s['running']}")
        print(f"  Success:        {s['success']}")
        print(f"  Failed:         {s['failed']}")
        print(f"  Skipped:        {s['skipped']}")
        print('=' * 60)

        # Show failed files
        failed = self.get_files_by_status('failed')
        if failed:
            print('\n  FAILED FILES:')
            for entry in failed:
                print(f"    - {entry['file_name']}: {entry.get('error', 'unknown error')}")
            print()


def get_s3_client():
    """Create S3 client for SeaweedFS."""
    endpoint = os.getenv('SEAWEEDFS_ENDPOINT')
    if not endpoint:
        host = os.getenv('SEAWEEDFS_HOST', 'localhost')
        port = os.getenv('SEAWEEDFS_PORT', '8333')
        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()
    return SUPPORTED_EXTENSIONS.get(ext, '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 collect_files(path: Path) -> list[Path]:
    """Collect all supported files from a directory recursively."""
    if path.is_file():
        return [path] if is_supported_file(path) else []

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

    return []


def upload_file(s3_client, file_path: Path, bucket: str, prefix: str) -> dict:
    """Upload a single file to SeaweedFS. Returns upload info."""
    object_key = f"{prefix}/{file_path.name}" if prefix else file_path.name
    object_key = object_key.lstrip('/')
    content_type = get_content_type(file_path)
    document_id = str(uuid.uuid4())

    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},
        )

    return {
        'document_id': document_id,
        'bucket': bucket,
        'object_key': object_key,
        'content_type': content_type,
    }


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 Airflow token: {response.status_code} - {response.text}")
            return None
    except Exception as e:
        print(f"  Failed to get Airflow token: {e}")
        return None


def trigger_dag(upload_info: dict, airflow_url: str, token: str) -> str | None:
    """Trigger the document_extraction_dag. Returns dag_run_id or None."""
    dag_id = 'document_extraction_dag'
    url = f"{airflow_url}/api/v2/dags/{dag_id}/dagRuns"

    payload = {
        'logical_date': None,
        'conf': {
            'document_id': upload_info['document_id'],
            'bucket': upload_info['bucket'],
            'object_key': upload_info['object_key'],
        },
    }

    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):
            return response.json().get('dag_run_id')
        else:
            print(f"  Failed to trigger DAG: {response.status_code} - {response.text}")
            return None
    except Exception as e:
        print(f"  Failed to trigger DAG: {e}")
        return None


def check_dag_run_status(dag_run_id: str, airflow_url: str, token: str) -> str | None:
    """Check status of a DAG run. Returns state string or None on error."""
    dag_id = 'document_extraction_dag'
    url = f"{airflow_url}/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}"

    try:
        response = requests.get(
            url,
            headers={'Authorization': f'Bearer {token}'},
            timeout=30,
        )
        if response.status_code == 200:
            return response.json().get('state')
        else:
            return None
    except Exception:
        return None


def phase_upload(s3_client, tracker: ProgressTracker, bucket: str, prefix: str, dry_run: bool):
    """Phase 1: Upload all pending files to SeaweedFS."""
    pending = tracker.get_files_by_status('pending')
    if not pending:
        print('  No pending files to upload.')
        return

    print(f'\n  Uploading {len(pending)} file(s)...')

    for entry in pending:
        file_path = Path(entry['file_path'])
        file_name = entry['file_name']

        if dry_run:
            print(f"  [DRY RUN] Would upload: {file_name}")
            tracker.update_status(file_name, 'uploaded',
                                  document_id=str(uuid.uuid4()),
                                  bucket=bucket,
                                  object_key=f"{prefix}/{file_name}".lstrip('/'))
            continue

        tracker.update_status(file_name, 'uploading')
        try:
            info = upload_file(s3_client, file_path, bucket, prefix)
            tracker.update_status(
                file_name, 'uploaded',
                document_id=info['document_id'],
                bucket=info['bucket'],
                object_key=info['object_key'],
                upload_time=datetime.now(timezone.utc).isoformat(),
            )
            size = entry['size_bytes']
            print(f"  Uploaded: {file_name} ({size:,} bytes)")
        except Exception as e:
            tracker.update_status(file_name, 'failed', error=f"Upload error: {e}")
            print(f"  FAILED:   {file_name} - {e}")


def phase_trigger(tracker: ProgressTracker, airflow_url: str, token: str, dry_run: bool):
    """Phase 2: Trigger DAG runs for all uploaded files."""
    uploaded = tracker.get_files_by_status('uploaded')
    if not uploaded:
        print('  No uploaded files to trigger.')
        return

    print(f'\n  Triggering DAG for {len(uploaded)} file(s)...')

    for entry in uploaded:
        file_name = entry['file_name']

        if dry_run:
            print(f"  [DRY RUN] Would trigger DAG for: {file_name}")
            tracker.update_status(file_name, 'success', dag_run_id='dry-run')
            continue

        upload_info = {
            'document_id': entry['document_id'],
            'bucket': entry['bucket'],
            'object_key': entry['object_key'],
        }
        dag_run_id = trigger_dag(upload_info, airflow_url, token)

        if dag_run_id:
            tracker.update_status(
                file_name, 'dag_triggered',
                dag_run_id=dag_run_id,
                trigger_time=datetime.now(timezone.utc).isoformat(),
            )
            print(f"  Triggered: {file_name} -> {dag_run_id}")
        else:
            tracker.update_status(file_name, 'failed', error='Failed to trigger DAG')
            print(f"  FAILED:    {file_name} - could not trigger DAG")


def phase_monitor(tracker: ProgressTracker, airflow_url: str, token: str, poll_interval: int):
    """Phase 3: Monitor DAG runs until all reach terminal state."""
    active = tracker.get_files_by_status('dag_triggered', 'running')
    if not active:
        print('  No active DAG runs to monitor.')
        return

    print(f'\n  Monitoring {len(active)} DAG run(s) (poll every {poll_interval}s)...')

    while True:
        active = tracker.get_files_by_status('dag_triggered', 'running')
        if not active:
            break

        for entry in active:
            file_name = entry['file_name']
            dag_run_id = entry['dag_run_id']

            state = check_dag_run_status(dag_run_id, airflow_url, token)

            if state is None:
                continue  # Transient error, retry next cycle

            if state in TERMINAL_STATES:
                tracker.update_status(
                    file_name, state,
                    completion_time=datetime.now(timezone.utc).isoformat(),
                    error='DAG run failed' if state == 'failed' else None,
                )
                icon = 'OK' if state == 'success' else 'FAIL'
                print(f"  [{icon}] {file_name} -> {state}")
            elif state in RUNNING_STATES:
                if entry['status'] != 'running':
                    tracker.update_status(file_name, 'running')

        # Recheck
        still_active = tracker.get_files_by_status('dag_triggered', 'running')
        if not still_active:
            break

        time.sleep(poll_interval)

    tracker.save()


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

    parser.add_argument(
        'path',
        type=Path,
        nargs='?',
        help='Folder or file 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(
        '--no-trigger',
        action='store_true',
        help='Upload only, do not trigger DAG',
    )
    parser.add_argument(
        '--dry-run', '-n',
        action='store_true',
        help='Show what would be done without actually doing it',
    )
    parser.add_argument(
        '--resume', '-r',
        type=str,
        default=None,
        help='Resume from a previous progress JSON file',
    )
    parser.add_argument(
        '--poll-interval',
        type=int,
        default=15,
        help='Seconds between DAG status checks (default: 15)',
    )
    parser.add_argument(
        '--progress-file',
        type=str,
        default=None,
        help='Custom path for the progress JSON file',
    )
    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()

    # --- Resume mode ---
    if args.resume:
        print(f"Resuming from: {args.resume}")
        tracker = ProgressTracker(args.resume)
        if not tracker.load():
            print(f"Error: cannot load progress file: {args.resume}")
            sys.exit(1)
        tracker.print_summary()
    else:
        # --- New upload mode ---
        if not args.path:
            parser.error('path is required unless --resume is used')

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

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

        # Create progress file
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        progress_file = args.progress_file or f"progress_{timestamp}.json"
        tracker = ProgressTracker(progress_file)

        for f in files:
            tracker.add_file(str(f), f.name, f.stat().st_size)
        tracker.save()

        print(f"Found {len(files)} supported file(s)")
        print(f"Progress file: {progress_file}")
        print(f"Bucket: {args.bucket}")
        if args.prefix:
            print(f"Prefix: {args.prefix}")

    # --- Phase 1: Upload ---
    s3_client = None
    if not args.dry_run:
        s3_client = get_s3_client()
        if not ensure_bucket_exists(s3_client, args.bucket):
            print("Failed to ensure bucket exists")
            sys.exit(1)

    phase_upload(s3_client, tracker, args.bucket, args.prefix, args.dry_run)
    tracker.print_summary()

    if args.no_trigger:
        print("\n--no-trigger set. Stopping after upload.")
        return

    # --- Phase 2: Trigger DAG ---
    token = None
    if not args.dry_run:
        token = get_airflow_token(args.airflow_url, args.airflow_user, args.airflow_password)
        if not token:
            print("Failed to get Airflow token. Files are uploaded but DAG not triggered.")
            print(f"Re-run with: --resume {tracker.progress_file}")
            sys.exit(1)

    phase_trigger(tracker, args.airflow_url, token, args.dry_run)
    tracker.print_summary()

    if args.dry_run:
        print("\n[DRY RUN] Complete.")
        return

    # --- Phase 3: Monitor ---
    print('\n--- Monitoring DAG runs (Ctrl+C to stop, resume later with --resume) ---')
    try:
        phase_monitor(tracker, args.airflow_url, token, args.poll_interval)
    except KeyboardInterrupt:
        print(f"\n\nInterrupted. Resume with: python {sys.argv[0]} --resume {tracker.progress_file}")
        tracker.save()
        sys.exit(0)

    tracker.print_summary()
    print(f"\nProgress saved to: {tracker.progress_file}")


if __name__ == '__main__':
    main()
