"""Sentrial CLI - Command line interface for running experiments"""

import argparse
import sys
import os
import importlib.util


def load_agent_module(agent_path: str):
    """Load a Python module from file path."""
    if not os.path.exists(agent_path):
        # Try relative to current directory
        if not os.path.exists(agent_path):
            raise FileNotFoundError(f"Agent file not found: {agent_path}")

    spec = importlib.util.spec_from_file_location("agent_module", agent_path)
    module = importlib.util.module_from_spec(spec)
    sys.modules["agent_module"] = module
    spec.loader.exec_module(module)
    return module


def run_experiment(experiment_id: str, agent_path: str = None, parallel: bool = False, max_workers: int = 4):
    """Run an experiment by ID."""
    from .client import SentrialClient
    from .experiment import Experiment
    from .context import set_experiment_context, clear_experiment_context

    api_key = os.environ.get("SENTRIAL_API_KEY")
    api_url = os.environ.get("SENTRIAL_API_URL", "http://localhost:3001")

    if not api_key:
        print("Error: SENTRIAL_API_KEY environment variable not set")
        sys.exit(1)

    print(f"Loading experiment {experiment_id}...")

    try:
        experiment = Experiment(
            experiment_id=experiment_id,
            api_key=api_key,
            api_url=api_url,
        )

        print(f"Experiment: {experiment.name}")
        print(f"Variants: {len(experiment.variants)}")
        print(f"Test cases: {len(experiment.test_cases)}")
        print()

        # Show variants
        for v in experiment.variants:
            prompt_preview = v.system_prompt[:50] + "..." if len(v.system_prompt) > 50 else v.system_prompt
            print(f"  - {v.name}: {prompt_preview}")
        print()

        if not agent_path:
            print("No agent specified. Use --agent to specify your agent file.")
            print()
            print("Example:")
            print(f"  sentrial experiment run {experiment_id} --agent my_agent.py")
            print()
            print("Your agent should use get_system_prompt() to receive experiment prompts:")
            print()
            print("    from sentrial import get_system_prompt")
            print()
            print("    def run_agent(query: str) -> str:")
            print("        prompt = get_system_prompt('Your default prompt')")
            print("        # ... rest of your agent code")
            print()
            return

        # Load the agent module
        print(f"Loading agent from: {agent_path}")

        # Clear any previously registered entry points before loading
        from .decorators import get_session_entry_points, clear_session_entry_points
        clear_session_entry_points()

        agent_module = load_agent_module(agent_path)

        # First, look for @session-decorated functions (preferred)
        entry_points = get_session_entry_points()
        agent_fn = None

        if entry_points:
            agent_fn = entry_points[0]  # Use the first @session-decorated function
            print(f"Found @session-decorated agent: {agent_fn.__name__}")
        else:
            # Fall back to name-based lookup for backwards compatibility
            for fn_name in ['run_agent', 'run_weather_agent', 'run_experiment_agent', 'agent']:
                if hasattr(agent_module, fn_name):
                    agent_fn = getattr(agent_module, fn_name)
                    print(f"Found agent function by name: {fn_name}")
                    break

        if not agent_fn:
            print("Error: Could not find agent function in module.")
            print()
            print("Option 1 (recommended): Use the @session decorator:")
            print("    from sentrial import session, get_system_prompt")
            print()
            print("    @session('my-agent')")
            print("    def my_agent(query: str):")
            print("        prompt = get_system_prompt('default')")
            print("        return response")
            print()
            print("Option 2: Name your function one of: run_agent, agent")
            sys.exit(1)

        # Create client for tracking
        client = SentrialClient(api_key=api_key, api_url=api_url)

        # Run experiment
        print()
        print("=" * 60)
        print("RUNNING EXPERIMENT")
        print("=" * 60)

        total_runs = len(experiment.variants) * len(experiment.test_cases)
        completed = 0
        results = []

        for variant in experiment.variants:
            for test_case in experiment.test_cases:
                completed += 1
                print(f"\n[{completed}/{total_runs}] Running: {variant.name} x {test_case.session_id[:8]}...")
                print(f"  Input: {test_case.user_input[:60]}..." if len(test_case.user_input) > 60 else f"  Input: {test_case.user_input}")
                print(f"  Prompt: {variant.system_prompt[:60]}..." if len(variant.system_prompt) > 60 else f"  Prompt: {variant.system_prompt}")

                with experiment.track_run(variant.name, test_case.session_id) as tracker:
                    try:
                        import time
                        start_time = time.time()

                        # Set experiment context (agent reads via get_system_prompt())
                        set_experiment_context(
                            system_prompt=variant.system_prompt,
                            variant_name=variant.name,
                            experiment_id=experiment_id,
                            test_case_id=test_case.session_id,
                        )

                        # Run the agent - just pass the input, nothing else required
                        result = agent_fn(test_case.user_input)

                        # Clear the context
                        clear_experiment_context()

                        duration_ms = int((time.time() - start_time) * 1000)

                        # Extract session_id if returned as tuple (optional - agent doesn't have to)
                        session_id = None
                        if isinstance(result, tuple) and len(result) >= 2:
                            output, session_id = result[0], result[1]
                        else:
                            output = result

                        # Link the agent's session to this experiment run
                        if session_id:
                            tracker.set_result_session_id(session_id)

                        print(f"  ✓ Success (took {duration_ms}ms)")
                        results.append({"variant": variant.name, "success": True, "duration_ms": duration_ms})

                    except Exception as e:
                        clear_experiment_context()  # Clean up on error
                        print(f"  ✗ Failed: {e}")
                        tracker.set_error(str(e))
                        results.append({"variant": variant.name, "success": False, "error": str(e)})

        print()
        print("=" * 60)
        print("EXPERIMENT COMPLETE")
        print("=" * 60)
        successful = sum(1 for r in results if r["success"])
        print(f"Results: {successful}/{total_runs} successful")
        print()
        print(f"View results at: {api_url.replace(':3001', ':3000')}/experiments/{experiment_id}")

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        prog="sentrial",
        description="Sentrial CLI - AI Agent Observability"
    )

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

    # experiment command
    exp_parser = subparsers.add_parser("experiment", help="Experiment commands")
    exp_subparsers = exp_parser.add_subparsers(dest="exp_command", help="Experiment subcommands")

    # experiment run
    run_parser = exp_subparsers.add_parser("run", help="Run an experiment")
    run_parser.add_argument("experiment_id", help="Experiment ID to run")
    run_parser.add_argument("--agent", "-a", help="Path to agent Python file")
    run_parser.add_argument("--parallel", action="store_true", help="Run variants in parallel")
    run_parser.add_argument("--max-workers", type=int, default=4, help="Max parallel workers")

    # experiment info
    info_parser = exp_subparsers.add_parser("info", help="Show experiment info")
    info_parser.add_argument("experiment_id", help="Experiment ID")

    args = parser.parse_args()

    if args.command == "experiment":
        if args.exp_command == "run":
            run_experiment(args.experiment_id, args.agent, args.parallel, args.max_workers)
        elif args.exp_command == "info":
            run_experiment(args.experiment_id)
        else:
            exp_parser.print_help()
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
