#!/usr/bin/env python3
"""Test script for the upload endpoint on Airflow webserver.

Usage:
    # Quick connectivity check (no file needed)
    python scripts/test_upload_endpoint.py --check

    # Upload a real file
    python scripts/test_upload_endpoint.py path/to/test.xlsx

    # Upload with custom host
    python scripts/test_upload_endpoint.py --host http://localhost:8081 path/to/test.xlsx

Requires HMAC_KEY_RING env var (or HMAC_KEY_RING_FILE) to be set.
"""

import argparse
import json
import sys
import uuid

from dotenv import load_dotenv
load_dotenv()

import requests
from textiq_hmac import load_key_ring, sign_context_headers

# Sample UUIDs matching tech-spec examples
DEFAULT_TENANT_ID = "28aa3971-48af-4f2a-aa73-9e4f89fbfba3"
DEFAULT_NODE_ID = "22222222-2222-2222-2222-222222222222"
DEFAULT_USER_ID = "33333333-3333-3333-3333-333333333333"
DEFAULT_ROLE = "content_specialist"


def sign_request(tenant_id, node_id, user_id, role):
    """Generate HMAC headers using textiq_hmac package."""
    key_ring = load_key_ring()
    headers = {
        "X-Tenant-ID": tenant_id,
        "X-Node-ID": node_id,
        "X-User-ID": user_id,
        "X-Role": role,
    }
    return sign_context_headers(headers, key_ring)


def check_endpoint(host):
    """Quick connectivity check — GET should return 405 if route exists."""
    url = f"{host}/api/v1/documents/upload"
    print(f"GET {url}")
    try:
        r = requests.get(url, timeout=5)
        print(f"  Status: {r.status_code}")
        if r.status_code in (405, 401):
            print(f"  -> Route exists ({r.status_code} = expected for unauthenticated GET)")
            return True
        elif r.status_code == 404:
            print("  -> Route NOT found. Plugin may not be loaded.")
            print("     Run: docker exec textiq-airflow-webserver airflow plugins")
            return False
        else:
            print(f"  -> Unexpected status. Body: {r.text[:200]}")
            return False
    except requests.ConnectionError as e:
        print(e)
        print(f"  -> Cannot connect to {host}. Is the webserver running?")
        return False


def test_no_auth(host):
    """POST with file but without auth headers — should return 401."""
    url = f"{host}/api/v1/documents/upload"
    print(f"\nPOST {url} (no auth headers, dummy file)")
    # Must include a file or FastAPI returns 422 before our handler runs
    import io
    dummy = io.BytesIO(b"PK dummy content")
    r = requests.post(url, files={"file": ("test.xlsx", dummy)}, timeout=10)
    print(f"  Status: {r.status_code}")
    print(f"  Body: {r.text[:200]}")
    if r.status_code == 401:
        print("  -> PASS: Got 401 as expected")
    else:
        print(f"  -> FAIL: Expected 401, got {r.status_code}")


def upload_file(host, filepath, tenant_id, node_id, user_id, role, correlation_id=None):
    """Upload a file with signed HMAC headers."""
    url = f"{host}/api/v1/documents/upload"
    headers = sign_request(tenant_id, node_id, user_id, role)

    correlation_id = correlation_id or str(uuid.uuid4())
    headers["X-Correlation-ID"] = correlation_id

    print(f"\nPOST {url}")
    print(f"  File: {filepath}")
    print(f"  Tenant: {tenant_id}")
    print(f"  Node: {node_id}")
    print(f"  Correlation-ID: {correlation_id}")
    print(f"  Timestamp: {headers['X-Timestamp']}")

    with open(filepath, "rb") as f:
        files = {"file": (filepath.split("/")[-1], f)}
        r = requests.post(url, headers=headers, files=files, timeout=30)

    print(f"  Status: {r.status_code}")
    print(f"  Response X-Correlation-ID: {r.headers.get('X-Correlation-ID', 'NOT PRESENT')}")
    print(f"  Response X-Request-ID: {r.headers.get('X-Request-ID', 'NOT PRESENT')}")
    try:
        body = r.json()
        print(f"  Body: {json.dumps(body, indent=2)}")
    except Exception:
        print(f"  Body: {r.text[:500]}")

    if r.status_code == 202:
        print("  -> PASS: Upload accepted")
        if isinstance(body, dict):
            print(f"  document_id: {body.get('document_id')}")
            print(f"  status: {body.get('status')}")
    else:
        print(f"  -> FAIL: Expected 202, got {r.status_code}")


def main():
    parser = argparse.ArgumentParser(description="Test DE upload endpoint")
    parser.add_argument("file", nargs="?", help="File to upload")
    parser.add_argument("--host", default="http://localhost:8081", help="Airflow webserver URL")
    parser.add_argument("--check", action="store_true", help="Quick connectivity check only")
    parser.add_argument("--tenant-id", default=DEFAULT_TENANT_ID)
    parser.add_argument("--node-id", default=DEFAULT_NODE_ID)
    parser.add_argument("--user-id", default=DEFAULT_USER_ID)
    parser.add_argument("--role", default=DEFAULT_ROLE)
    parser.add_argument("--correlation-id", default=None, help="X-Correlation-ID (auto-generated if omitted)")
    args = parser.parse_args()

    print(f"=== Upload Endpoint Test ({args.host}) ===\n")

    # Always run connectivity check
    alive = check_endpoint(args.host)
    if not alive:
        sys.exit(1)

    if args.check:
        test_no_auth(args.host)
        return

    if not args.file:
        print("\nNo file specified. Running auth-only tests.\n")
        test_no_auth(args.host)
        return

    upload_file(
        args.host,
        args.file,
        args.tenant_id,
        args.node_id,
        args.user_id,
        args.role,
        args.correlation_id,
    )


if __name__ == "__main__":
    main()