#!/usr/bin/env python3
"""mca-instrument CLI - Zero-code OpenTelemetry instrumentation wrapper.

This command wraps opentelemetry-instrument to enable full OTel auto-instrumentation
of model scripts without requiring code changes.

Usage:
    mca-instrument --model-id mdl-001 --team clinical-ai -- python my_model.py

Requirements:
    pip install "mca-sdk[instrument]"
    opentelemetry-bootstrap -a install  # Install auto-instrumentation libraries
"""

import argparse
import logging
import os
import shutil
import signal
import subprocess
import sys
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


def build_parser() -> argparse.ArgumentParser:
    """Build argument parser for mca-instrument command."""
    parser = argparse.ArgumentParser(
        description="Instrument model scripts with OpenTelemetry (zero code changes)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Basic usage (use -- to separate wrapper args from command)
  mca-instrument --model-id mdl-001 --team clinical-ai -- python my_model.py

  # With all options
  mca-instrument \\
    --model-id mdl-001 \\
    --team clinical-ai \\
    --service-name my-service \\
    --collector-endpoint http://localhost:4318 \\
    --protocol http/protobuf \\
    -- python my_model.py --arg value

  # Using environment variables (higher priority)
  export MCA_MODEL_ID=mdl-001
  export MCA_TEAM_NAME=clinical-ai
  mca-instrument -- python my_model.py

Requirements:
  pip install "mca-sdk[instrument]"
  opentelemetry-bootstrap -a install
""",
    )

    # MCA SDK configuration
    parser.add_argument(
        "--model-id",
        help="Model identifier (or set MCA_MODEL_ID env var)",
    )
    parser.add_argument(
        "--team",
        dest="team_name",
        help="Team name (or set MCA_TEAM_NAME env var)",
    )
    parser.add_argument(
        "--service-name",
        help="Service name (defaults to model-id if not set)",
    )
    parser.add_argument(
        "--collector-endpoint",
        help="OpenTelemetry Collector endpoint (or set MCA_COLLECTOR_ENDPOINT)",
    )
    parser.add_argument(
        "--protocol",
        choices=["http/protobuf", "grpc"],
        help="OTLP protocol (sets OTEL_EXPORTER_OTLP_PROTOCOL)",
    )
    parser.add_argument(
        "--registry-url",
        help="Model Registry API URL (or set MCA_REGISTRY_URL)",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug logging",
    )

    return parser


def inject_env_vars(
    args: argparse.Namespace, env: Optional[Dict[str, str]] = None
) -> Dict[str, str]:
    """Build environment dict with MCA_* variables.

    CRITICAL: Never override existing environment variables (AC 3).
    CLI args are lower priority than existing env vars.

    Remediation Fix: Calculate derived variables AFTER merging CLI and environment.

    Args:
        args: Parsed command-line arguments
        env: Base environment (defaults to os.environ.copy())

    Returns:
        Environment dict with MCA_* variables injected
    """
    if env is None:
        env = os.environ.copy()

    # Step 1: Determine final values (env vars take precedence over CLI args)
    final_model_id = env.get("MCA_MODEL_ID") or args.model_id
    final_team_name = env.get("MCA_TEAM_NAME") or args.team_name
    final_service_name = env.get("MCA_SERVICE_NAME") or args.service_name or final_model_id
    final_collector_endpoint = env.get("MCA_COLLECTOR_ENDPOINT") or args.collector_endpoint
    final_registry_url = env.get("MCA_REGISTRY_URL") or args.registry_url
    final_protocol = env.get("OTEL_EXPORTER_OTLP_PROTOCOL") or args.protocol or "http/protobuf"

    # Step 2: Set variables (only if we have a value)
    if final_model_id:
        env["MCA_MODEL_ID"] = final_model_id
    if final_team_name:
        env["MCA_TEAM_NAME"] = final_team_name
    if final_service_name:
        env["MCA_SERVICE_NAME"] = final_service_name
        # Also set OTEL_SERVICE_NAME for opentelemetry-instrument
        if "OTEL_SERVICE_NAME" not in env:
            env["OTEL_SERVICE_NAME"] = final_service_name
    if final_collector_endpoint:
        env["MCA_COLLECTOR_ENDPOINT"] = final_collector_endpoint
        # Map to OpenTelemetry environment variable
        if "OTEL_EXPORTER_OTLP_ENDPOINT" not in env:
            env["OTEL_EXPORTER_OTLP_ENDPOINT"] = final_collector_endpoint
    if final_registry_url:
        env["MCA_REGISTRY_URL"] = final_registry_url
    if final_protocol:
        env["OTEL_EXPORTER_OTLP_PROTOCOL"] = final_protocol
        # Set specific component protocol overrides to prevent OpenTelemetry from falling back
        # to its default 'grpc' for specific signals if they aren't explicitly configured.
        if "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL" not in env:
            env["OTEL_EXPORTER_OTLP_TRACES_PROTOCOL"] = final_protocol
        if "OTEL_EXPORTER_OTLP_METRICS_PROTOCOL" not in env:
            env["OTEL_EXPORTER_OTLP_METRICS_PROTOCOL"] = final_protocol
        if "OTEL_EXPORTER_OTLP_LOGS_PROTOCOL" not in env:
            env["OTEL_EXPORTER_OTLP_LOGS_PROTOCOL"] = final_protocol

    # Debug mode
    if args.debug and "MCA_DEBUG_MODE" not in env:
        env["MCA_DEBUG_MODE"] = "true"

    return env


def find_otel_instrument() -> Optional[str]:
    """Find opentelemetry-instrument command on PATH.

    Remediation Fix: Only look for the binary on PATH (user's target environment).
    Do NOT fallback to sys.executable -m pattern (wrong virtual environment).

    Returns:
        Path to opentelemetry-instrument binary, or None if not found
    """
    # Only look on PATH - this respects user's target environment
    path = shutil.which("opentelemetry-instrument")
    return path


def run_instrumented(
    otel_cmd: str, script_cmd: List[str], env: Dict[str, str], timeout: float = 5.0
) -> int:
    """Run instrumented command with OpenTelemetry and MCA Bootstrap.

    Remediation Fix: Proper signal handling with timeout and SIGKILL escalation.
    Returns POSIX-compliant exit codes for signals (128 + signal number).

    Crucial fix: We must inject MCA _bootstrap.py so MCAClient initializes
    and sets up GCP Exporters, otherwise it's just raw OpenTelemetry.

    Args:
        otel_cmd: Path to opentelemetry-instrument binary
        script_cmd: User's script command
        env: Environment dict
        timeout: Seconds to wait for graceful shutdown before SIGKILL

    Returns:
        Exit code from subprocess (POSIX-compliant)
    """
    # Build full command: opentelemetry-instrument [script command]
    full_cmd = [otel_cmd] + script_cmd

    temp_dir = None
    try:
        # Setup bootstrap injection if command is Python
        if script_cmd and ("python" in script_cmd[0].lower() or script_cmd[0].endswith(".py")):
            from mca_sdk.cli.run import setup_bootstrap_injection

            temp_dir = setup_bootstrap_injection(env)

        # Spawn subprocess
        process = subprocess.Popen(full_cmd, env=env)

        # Setup signal forwarding (POSIX only)
        def forward_signal(signum, frame):
            """Forward signal to child, wait with timeout, then SIGKILL if necessary."""
            if process.poll() is None:  # Child still running
                process.send_signal(signum)
                try:
                    process.wait(timeout=timeout)
                except subprocess.TimeoutExpired:
                    # Child didn't exit gracefully, force kill
                    logger.warning(
                        "Child process did not terminate after %s seconds, sending SIGKILL",
                        timeout,
                    )
                    process.kill()
                    process.wait()

        if hasattr(signal, "SIGTERM"):
            signal.signal(signal.SIGTERM, forward_signal)
        if hasattr(signal, "SIGINT"):
            signal.signal(signal.SIGINT, forward_signal)

        # Wait for completion
        process.wait()

        # Handle exit codes properly
        returncode = process.returncode
        if returncode < 0:
            # Process was terminated by signal
            # Return POSIX-compliant exit code: 128 + signal number
            return 128 + abs(returncode)
        return returncode

    except FileNotFoundError:
        print(f"Error: Command not found: {full_cmd[0]}", file=sys.stderr)
        return 127
    except Exception as e:
        print(f"Error running command: {e}", file=sys.stderr)
        return 1
    finally:
        # Cleanup injected files
        if temp_dir and os.path.exists(temp_dir):
            try:
                shutil.rmtree(temp_dir)
            except Exception as e:
                logger.debug("Failed to clean up temp dir %s: %s", temp_dir, e)


def validate_config(env: Dict[str, str]) -> bool:
    """Validate required configuration.

    Args:
        env: Environment dict

    Returns:
        True if valid, False otherwise (prints error to stderr)
    """
    required = {
        "MCA_MODEL_ID": "model ID (use --model-id or set MCA_MODEL_ID)",
        "MCA_TEAM_NAME": "team name (use --team or set MCA_TEAM_NAME)",
    }

    missing = []
    for key, desc in required.items():
        if key not in env or not env[key]:
            missing.append(desc)

    if missing:
        print(
            f"Error: Missing required configuration: {', '.join(missing)}",
            file=sys.stderr,
        )
        return False

    return True


def main() -> int:
    """Main entry point for mca-instrument CLI."""
    # Remediation Fix: Use parse_known_args() for robust argument parsing
    parser = build_parser()

    # Parse known args, everything after -- goes to command
    if "--" in sys.argv:
        sep_index = sys.argv.index("--")
        wrapper_args = sys.argv[1:sep_index]
        command = sys.argv[sep_index + 1 :]
    else:
        # No separator, try to parse what we can
        wrapper_args = []
        command = []
        for arg in sys.argv[1:]:
            if arg.startswith("-") and not command:
                wrapper_args.append(arg)
            else:
                command.append(arg)
                # Once we hit the command, everything else is part of it
                command.extend(sys.argv[sys.argv.index(arg) + 1 :])
                break

    # Parse wrapper arguments
    args = parser.parse_args(wrapper_args)

    # Validate command
    if not command:
        parser.print_help()
        print("\nError: No command specified", file=sys.stderr)
        return 2

    # Inject environment variables
    env = inject_env_vars(args)

    # Validate configuration (AC 8)
    if not validate_config(env):
        return 2

    # Find OpenTelemetry instrumentation command
    otel_cmd = find_otel_instrument()
    if not otel_cmd:
        print(
            "Error: opentelemetry-instrument not found on PATH.\n"
            "\n"
            "To fix this:\n"
            "  1. pip install 'mca-sdk[instrument]'\n"
            "  2. opentelemetry-bootstrap -a install\n"
            "\n"
            "Alternatively, use 'mca-run' which doesn't require OTel dependencies.",
            file=sys.stderr,
        )
        return 1

    # Run instrumented command (AC 5)
    return run_instrumented(otel_cmd, command, env)


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