#!/usr/bin/env python3
"""
Setup Airflow Connections for TextIQ Services.

This script configures Airflow connections for all TextIQ backend services.
Run this after Airflow is initialized to set up connections.

Usage:
    # From within the airflow-cli container:
    docker compose -f docker/docker-compose.airflow.yml --profile cli run --rm airflow-cli python /opt/airflow/src/../scripts/setup_airflow_connections.py

    # Or directly with airflow CLI:
    docker compose -f docker/docker-compose.airflow.yml exec airflow-webserver python /opt/airflow/src/../scripts/setup_airflow_connections.py
"""

import os
import subprocess
import sys
from typing import NamedTuple


class ConnectionConfig(NamedTuple):
    """Configuration for an Airflow connection."""

    conn_id: str
    conn_type: str
    host: str
    port: int | None = None
    login: str | None = None
    password: str | None = None
    schema: str | None = None
    extra: str | None = None
    description: str = ""


def get_env(key: str, default: str = "") -> str:
    """Get environment variable with default."""
    return os.environ.get(key, default)


def get_connections() -> list[ConnectionConfig]:
    """Get list of connections to configure."""
    return [
        # PostgreSQL - TextIQ application database
        ConnectionConfig(
            conn_id="postgres_default",
            conn_type="postgres",
            host=get_env("TextIQ_POSTGRES_HOST", "postgres"),
            port=int(get_env("TextIQ_POSTGRES_PORT", "5432")),
            login=get_env("TextIQ_POSTGRES_USER", "textiq"),
            password=get_env("TextIQ_POSTGRES_PASSWORD", "textiq_dev_password"),
            schema=get_env("TextIQ_POSTGRES_DB", "textiq"),
            description="TextIQ PostgreSQL database for documents and records",
        ),
        # Milvus Vector Database
        ConnectionConfig(
            conn_id="milvus_default",
            conn_type="generic",
            host=get_env("TextIQ_MILVUS_HOST", "milvus"),
            port=int(get_env("TextIQ_MILVUS_PORT", "19530")),
            extra='{"protocol": "grpc"}',
            description="Milvus vector database for embeddings",
        ),
        # Redis
        ConnectionConfig(
            conn_id="redis_default",
            conn_type="redis",
            host=get_env("TextIQ_REDIS_HOST", "redis"),
            port=int(get_env("TextIQ_REDIS_PORT", "6379")),
            description="Redis for caching and task queue",
        ),
        # SeaweedFS Object Storage (S3-compatible)
        ConnectionConfig(
            conn_id="seaweedfs_default",
            conn_type="aws",
            host=get_env("TextIQ_SEAWEEDFS_HOST", "seaweedfs"),
            port=int(get_env("TextIQ_SEAWEEDFS_PORT", "8333")),
            login=get_env("TextIQ_SEAWEEDFS_ACCESS_KEY", ""),
            password=get_env("TextIQ_SEAWEEDFS_SECRET_KEY", ""),
            extra=(
                '{"endpoint_url": "http://'
                + get_env("TextIQ_SEAWEEDFS_HOST", "seaweedfs")
                + ":"
                + get_env("TextIQ_SEAWEEDFS_PORT", "8333")
                + '"}'
            ),
            description="SeaweedFS S3-compatible object storage for documents",
        ),
        # Azure OpenAI
        ConnectionConfig(
            conn_id="azure_openai_default",
            conn_type="http",
            host=get_env(
                "AZURE_OPENAI_ENDPOINT", "https://your-resource.openai.azure.com/"
            ),
            password=get_env("AZURE_OPENAI_API_KEY", ""),
            extra=(
                '{"api_version": "'
                + get_env("AZURE_OPENAI_API_VERSION", "2024-02-15-preview")
                + '", "deployment": "'
                + get_env("AZURE_OPENAI_DEPLOYMENT", "gpt-4o")
                + '", "embedding_deployment": "'
                + get_env("AZURE_OPENAI_EMBEDDING_DEPLOYMENT", "text-embedding-3-large")
                + '"}'
            ),
            description="Azure OpenAI for embeddings and LLM",
        ),
    ]


def create_connection(conn: ConnectionConfig) -> bool:
    """Create or update an Airflow connection using CLI."""
    cmd = [
        "airflow",
        "connections",
        "add",
        conn.conn_id,
        "--conn-type",
        conn.conn_type,
    ]

    if conn.host:
        cmd.extend(["--conn-host", conn.host])
    if conn.port:
        cmd.extend(["--conn-port", str(conn.port)])
    if conn.login:
        cmd.extend(["--conn-login", conn.login])
    if conn.password:
        cmd.extend(["--conn-password", conn.password])
    if conn.schema:
        cmd.extend(["--conn-schema", conn.schema])
    if conn.extra:
        cmd.extend(["--conn-extra", conn.extra])
    if conn.description:
        cmd.extend(["--conn-description", conn.description])

    # Try to delete existing connection first (ignore errors)
    subprocess.run(
        ["airflow", "connections", "delete", conn.conn_id],
        capture_output=True,
        check=False,
    )

    # Create the connection
    result = subprocess.run(cmd, capture_output=True, text=True, check=False)

    if result.returncode != 0:
        print(f"  ERROR: {result.stderr.strip()}")
        return False
    return True


def main() -> int:
    """Set up all Airflow connections."""
    print("=" * 60)
    print("TextIQ Airflow Connection Setup")
    print("=" * 60)

    connections = get_connections()
    success_count = 0
    failed_count = 0

    for conn in connections:
        print(f"\nConfiguring: {conn.conn_id}")
        print(f"  Type: {conn.conn_type}")
        print(f"  Host: {conn.host}:{conn.port}")

        if create_connection(conn):
            print("  Status: SUCCESS")
            success_count += 1
        else:
            print("  Status: FAILED")
            failed_count += 1

    print("\n" + "=" * 60)
    print(f"Setup Complete: {success_count} succeeded, {failed_count} failed")
    print("=" * 60)

    # List all connections
    print("\nCurrent Airflow connections:")
    subprocess.run(["airflow", "connections", "list"], check=False)

    return 0 if failed_count == 0 else 1


if __name__ == "__main__":
    sys.exit(main())