"""
Command-line interface for Auror.

This module provides CLI commands for running Auror workers and managing
the distributed inference system.
"""

import argparse
import logging
import sys
from typing import Optional

from auror.config import AurorConfig, get_default_config
from auror.core.messaging import MessagingCore
from auror.workers.ltxv import LTXVWorker
from auror.workers.ltxv_mock import LTXVWorkerMock
from auror.workers.sd35 import SD35Worker


def setup_logging(verbose: bool = False) -> None:
    """
    Configure logging for the application.

    Args:
        verbose: If True, set log level to DEBUG; otherwise INFO
    """
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def run_ltxv(config: AurorConfig, verbose: bool = False) -> None:
    """
    Run the LTXV worker.

    Args:
        config: Auror configuration
        verbose: Enable verbose logging
    """
    setup_logging(verbose)
    logger = logging.getLogger(__name__)

    logger.info("Starting Auror LTXV worker")
    logger.info(f"MQTT Broker: {config.mqtt_host}:{config.mqtt_port}")
    logger.info(f"MQTT Topic: {config.mqtt_topic}")
    logger.info(f"Pipeline Config: {config.pipeline_config}")
    logger.info(f"Mock Mode: {config.mock_mode}")

    # Create worker based on mock mode
    if config.mock_mode:
        worker = LTXVWorkerMock()
    else:
        worker = LTXVWorker(pipeline_config=config.pipeline_config)

    # Create messaging core
    core = MessagingCore(
        host=config.mqtt_host,
        port=config.mqtt_port,
        client_id=config.client_id,
        client_secret=config.client_secret,
        audience=config.audience,
        topic=config.mqtt_topic,
        auth0_token_url=config.auth0_token_url,
    )

    # Start processing
    try:
        core.start(worker)
    except KeyboardInterrupt:
        logger.info("Shutting down gracefully...")
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        sys.exit(1)


def run_sd35(
    config: AurorConfig,
    verbose: bool = False,
    model_id: str = "stabilityai/stable-diffusion-3.5-large",
    device: str = "cuda",
) -> None:
    """
    Run the SD35 worker.

    Args:
        config: Auror configuration
        verbose: Enable verbose logging
        model_id: HuggingFace model identifier for SD 3.5
        device: Device to run inference on (cuda or cpu)
    """
    setup_logging(verbose)
    logger = logging.getLogger(__name__)

    logger.info("Starting Auror SD35 worker")
    logger.info(f"MQTT Broker: {config.mqtt_host}:{config.mqtt_port}")
    logger.info(f"MQTT Topic: {config.mqtt_topic}")
    logger.info(f"Model ID: {model_id}")
    logger.info(f"Device: {device}")

    # Create SD35 worker
    worker = SD35Worker(model_id=model_id, device=device)

    # Create messaging core
    core = MessagingCore(
        host=config.mqtt_host,
        port=config.mqtt_port,
        client_id=config.client_id,
        client_secret=config.client_secret,
        audience=config.audience,
        topic=config.mqtt_topic,
        auth0_token_url=config.auth0_token_url,
    )

    # Start processing
    try:
        core.start(worker)
    except KeyboardInterrupt:
        logger.info("Shutting down gracefully...")
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        sys.exit(1)


def main() -> None:
    """
    Main entry point for the Auror CLI.
    """
    parser = argparse.ArgumentParser(
        prog="auror",
        description="Auror - Distributed AI inference task processing",
    )

    parser.add_argument(
        "--version",
        action="version",
        version="%(prog)s 0.1.0",
    )

    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Enable verbose logging (DEBUG level)",
    )

    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # LTXV worker command
    ltxv_parser = subparsers.add_parser(
        "ltxv",
        help="Run the LTX Video worker",
    )

    ltxv_parser.add_argument(
        "--host",
        type=str,
        help="MQTT broker hostname (overrides AUROR_MQTT_HOST)",
    )

    ltxv_parser.add_argument(
        "--port",
        type=int,
        help="MQTT broker port (overrides AUROR_MQTT_PORT)",
    )

    ltxv_parser.add_argument(
        "--topic",
        type=str,
        help="MQTT topic to subscribe to (overrides AUROR_MQTT_TOPIC)",
    )

    ltxv_parser.add_argument(
        "--client-id",
        type=str,
        help="OAuth2 client ID (overrides AUROR_CLIENT_ID)",
    )

    ltxv_parser.add_argument(
        "--client-secret",
        type=str,
        help="OAuth2 client secret (overrides AUROR_CLIENT_SECRET)",
    )

    ltxv_parser.add_argument(
        "--audience",
        type=str,
        help="OAuth2 audience (overrides AUROR_AUDIENCE)",
    )

    ltxv_parser.add_argument(
        "--pipeline-config",
        type=str,
        help="Path to LTX pipeline config (overrides AUROR_PIPELINE_CONFIG)",
    )

    ltxv_parser.add_argument(
        "--mock",
        action="store_true",
        help="Enable mock mode (simulate inference, overrides AUROR_MOCK_MODE)",
    )

    ltxv_parser.add_argument(
        "--use-defaults",
        action="store_true",
        help="Use default configuration (for development only)",
    )

    # SD35 worker command
    sd35_parser = subparsers.add_parser(
        "sd35",
        help="Run the Stable Diffusion 3.5 worker",
    )

    sd35_parser.add_argument(
        "--host",
        type=str,
        help="MQTT broker hostname (overrides AUROR_MQTT_HOST)",
    )

    sd35_parser.add_argument(
        "--port",
        type=int,
        help="MQTT broker port (overrides AUROR_MQTT_PORT)",
    )

    sd35_parser.add_argument(
        "--topic",
        type=str,
        help="MQTT topic to subscribe to (overrides AUROR_MQTT_TOPIC)",
    )

    sd35_parser.add_argument(
        "--client-id",
        type=str,
        help="OAuth2 client ID (overrides AUROR_CLIENT_ID)",
    )

    sd35_parser.add_argument(
        "--client-secret",
        type=str,
        help="OAuth2 client secret (overrides AUROR_CLIENT_SECRET)",
    )

    sd35_parser.add_argument(
        "--audience",
        type=str,
        help="OAuth2 audience (overrides AUROR_AUDIENCE)",
    )

    sd35_parser.add_argument(
        "--model-id",
        type=str,
        default="stabilityai/stable-diffusion-3.5-large",
        help="HuggingFace model identifier for SD 3.5 (default: stabilityai/stable-diffusion-3-5)",
    )

    sd35_parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        choices=["cuda", "cpu"],
        help="Device to run inference on (default: cuda)",
    )

    sd35_parser.add_argument(
        "--use-defaults",
        action="store_true",
        help="Use default configuration (for development only)",
    )

    args = parser.parse_args()

    # Show help if no command specified
    if not args.command:
        parser.print_help()
        sys.exit(0)

    # Handle LTXV command
    if args.command == "ltxv":
        # Get configuration
        if args.use_defaults:
            config = get_default_config()
        else:
            try:
                config = AurorConfig.from_env()
            except ValueError as e:
                print(f"Error: {e}", file=sys.stderr)
                print(
                    "\nYou can either:",
                    file=sys.stderr,
                )
                print("1. Set the required environment variables", file=sys.stderr)
                print(
                    "2. Use --use-defaults flag (development only)",
                    file=sys.stderr,
                )
                print(
                    "3. Override specific values with command-line arguments",
                    file=sys.stderr,
                )
                sys.exit(1)

        # Override with command-line arguments
        if args.host:
            config.mqtt_host = args.host
        if args.port:
            config.mqtt_port = args.port
        if args.topic:
            config.mqtt_topic = args.topic
        if args.client_id:
            config.client_id = args.client_id
        if args.client_secret:
            config.client_secret = args.client_secret
        if args.audience:
            config.audience = args.audience
        if args.pipeline_config:
            config.pipeline_config = args.pipeline_config
        if args.mock:
            config.mock_mode = True

        # Run worker
        run_ltxv(config, verbose=args.verbose)

    # Handle SD35 command
    elif args.command == "sd35":
        # Get configuration
        if args.use_defaults:
            config = get_default_config()
        else:
            try:
                config = AurorConfig.from_env()
            except ValueError as e:
                print(f"Error: {e}", file=sys.stderr)
                print(
                    "\nYou can either:",
                    file=sys.stderr,
                )
                print("1. Set the required environment variables", file=sys.stderr)
                print(
                    "2. Use --use-defaults flag (development only)",
                    file=sys.stderr,
                )
                print(
                    "3. Override specific values with command-line arguments",
                    file=sys.stderr,
                )
                sys.exit(1)

        # Override with command-line arguments
        if args.host:
            config.mqtt_host = args.host
        if args.port:
            config.mqtt_port = args.port
        if args.topic:
            config.mqtt_topic = args.topic
        if args.client_id:
            config.client_id = args.client_id
        if args.client_secret:
            config.client_secret = args.client_secret
        if args.audience:
            config.audience = args.audience

        # Run worker
        run_sd35(
            config,
            verbose=args.verbose,
            model_id=args.model_id,
            device=args.device,
        )


if __name__ == "__main__":
    main()
