#!/usr/bin/env python3
"""mca-run CLI - Zero-code MCAClient initialization wrapper.

This command runs model scripts with MCA_* environment variables injected,
enabling auto-initialization of MCAClient without requiring code changes.

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

Difference from mca-instrument:
    - No dependency on opentelemetry-instrumentation package
    - Injects MCA_* env vars and auto-initializes MCAClient
    - Uses PYTHONPATH injection for bootstrap code
"""

import argparse
import logging
import os
import signal
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


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

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

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

Note:
  This command does NOT require opentelemetry-instrumentation.
  For full auto-instrumentation, use mca-instrument instead.
""",
    )

    # MCA SDK configuration (same as instrument.py)
    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(
        "--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.
    CLI args are lower priority than existing env vars.

    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

    # 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
    if final_collector_endpoint:
        env["MCA_COLLECTOR_ENDPOINT"] = final_collector_endpoint
    if final_registry_url:
        env["MCA_REGISTRY_URL"] = final_registry_url

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

    return env


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 setup_bootstrap_injection(env: Dict[str, str]) -> str:
    """Setup bootstrap injection via sitecustomize.py in temp directory.

    Remediation Fix: Properly inject _bootstrap.py into child process via PYTHONPATH.

    Args:
        env: Environment dict to modify

    Returns:
        Path to temporary directory (caller must clean up)
    """
    # Create temporary directory for sitecustomize.py
    temp_dir = tempfile.mkdtemp(prefix="mca-run-")

    # Write sitecustomize.py that imports bootstrap
    sitecustomize_path = Path(temp_dir) / "sitecustomize.py"
    sitecustomize_path.write_text("""# MCA SDK auto-initialization via mca-run
import sys
import os

# 1. Run MCA Bootstrap
if os.environ.get('MCA_MODEL_ID'):
    try:
        import mca_sdk.cli._bootstrap
    except Exception as e:
        import warnings
        warnings.warn(f"MCA SDK bootstrap failed: {e}", stacklevel=2)

# 2. Support OpenTelemetry sitecustomize chaining
# opentelemetry-instrument uses sitecustomize to load its auto-instrumentation.
# Since Python only loads the FIRST sitecustomize.py it finds in PYTHONPATH,
# we must manually find and execute OpenTelemetry's sitecustomize if we are
# running under opentelemetry-instrument.
try:
    # Remove our temp directory from sys.path so we don't import ourselves infinitely
    current_dir = os.path.dirname(os.path.abspath(__file__))
    if current_dir in sys.path:
        sys.path.remove(current_dir)
        
    import importlib.util
    import site
    
    # Check if opentelemetry auto-instrumentation is trying to run
    if "opentelemetry.instrumentation.auto_instrumentation" in sys.path[0] or any("opentelemetry" in p for p in sys.path):
        import opentelemetry.instrumentation.auto_instrumentation.sitecustomize
except Exception as e:
    pass

""")

    # Prepend temp dir to PYTHONPATH
    current_pythonpath = env.get("PYTHONPATH", "")
    if current_pythonpath:
        env["PYTHONPATH"] = f"{temp_dir}:{current_pythonpath}"
    else:
        env["PYTHONPATH"] = temp_dir

    return temp_dir


def run_command(script_cmd: List[str], env: Dict[str, str], timeout: float = 5.0) -> int:
    """Run command with MCA_* environment variables and bootstrap injection.

    Remediation Fix:
    - Removed hardcoded sys.executable - execute command as-is
    - Added proper signal handling with timeout
    - Returns POSIX-compliant exit codes

    Args:
        script_cmd: User's script command (executed as-is, no sys.executable prepending)
        env: Environment dict
        timeout: Seconds to wait for graceful shutdown before SIGKILL

    Returns:
        Exit code from subprocess (POSIX-compliant)
    """
    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")):
            temp_dir = setup_bootstrap_injection(env)

        # Spawn subprocess (execute command exactly as provided)
        process = subprocess.Popen(script_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:
                    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: {script_cmd[0]}", file=sys.stderr)
        return 127
    except Exception as e:
        print(f"Error running command: {e}", file=sys.stderr)
        return 1
    finally:
        # Cleanup temp directory
        if temp_dir:
            try:
                import shutil

                shutil.rmtree(temp_dir, ignore_errors=True)
            except Exception:
                pass  # Best effort cleanup


def main() -> int:
    """Main entry point for mca-run CLI."""
    # Use same robust parsing as instrument.py
    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
    if not validate_config(env):
        return 2

    # Run command
    return run_command(command, env)


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