"""
Command-line interface for SkyDiscover
"""

import argparse
import asyncio
import logging
import multiprocessing
import os
import sys
import time

try:
    multiprocessing.set_start_method('spawn')
except RuntimeError:
    pass  

from skydiscover import MainController
from skydiscover.config import load_config, _parse_model_spec, _resolve_api_key_from_env, LLMModelConfig

logger = logging.getLogger(__name__)


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments"""
    parser = argparse.ArgumentParser(description="SkyDiscover - Flexible AI-Driven Scientific and Algorithmic Discovery")

    parser.add_argument("initial_program", help="Path to the initial program file")

    parser.add_argument(
        "evaluation_file", help="Path to the evaluation file containing an 'evaluate' function"
    )

    parser.add_argument("--config", "-c", help="Path to configuration file (YAML)", default=None)

    parser.add_argument("--output", "-o", help="Output directory for results", default=None)

    parser.add_argument(
        "--iterations", "-i", help="Maximum number of iterations", type=int, default=None
    )

    parser.add_argument(
        "--log-level",
        "-l",
        help="Logging level",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        default=None,
    )

    parser.add_argument(
        "--checkpoint",
        help="Path to checkpoint directory to resume from (e.g., outputs/checkpoints/checkpoint_50)",
        default=None,
    )

    parser.add_argument("--api-base", help="Base URL for the LLM API", default=None)

    parser.add_argument(
        "--codebase",
        help="Path to codebase directory for agentic mode. "
             "When provided, enables agentic generation so the LLM can "
             "read files and search the codebase before writing code.",
        default=None,
    )

    parser.add_argument(
        "--model", "-m",
        help="LLM model name (overrides config). Replaces the entire models list with a single model.",
        default=None,
    )

    parser.add_argument(
        "--search", "-s",
        help="Search algorithm to use",
        choices=["evox", "adaevolve", "best_of_n", "beam_search", "topk", "openevolve", "shinkaevolve", "gepa"],
        default=None,
    )

    return parser.parse_args()


async def main_async() -> int:
    """
    Main asynchronous entry point

    Returns:
        Exit code
    """
    args = parse_args()

    # Configure root logger using the same concise console format as logging_utils
    from skydiscover.search.utils.logging_utils import _ConsoleFormatter, _ConsoleFilter
    log_level = getattr(logging, args.log_level) if args.log_level else logging.WARNING
    root = logging.getLogger()
    root.setLevel(log_level)
    if not root.handlers:
        ch = logging.StreamHandler()
        ch.setFormatter(_ConsoleFormatter())
        ch.addFilter(_ConsoleFilter())
        root.addHandler(ch)
    # Always let our own logs through at INFO
    logging.getLogger("skydiscover").setLevel(logging.INFO)

    # Check if files exist
    if not os.path.exists(args.initial_program):
        print(f"Error: Initial program file '{args.initial_program}' not found")
        return 1

    if not os.path.exists(args.evaluation_file):
        print(f"Error: Evaluation file '{args.evaluation_file}' not found")
        return 1

    # Create config object with command-line overrides
    config = None
    has_cli_overrides = any([
        args.api_base, args.model, args.codebase, args.search,
    ])
    if args.config or has_cli_overrides:
        config = load_config(args.config)

        parsed_env_vars = None
        if args.model:
            provider, model_name, default_api_base, parsed_env_vars = _parse_model_spec(args.model)
            print(f"Resolved provider: {provider}")

            # api_base: explicit --api-base wins → provider default → error
            if args.api_base:
                config.llm.api_base = args.api_base
            elif default_api_base is not None:
                config.llm.api_base = default_api_base
            else:
                print(
                    f"Error: Provider '{provider}' requires --api-base.\n"
                    f"Example: --model {args.model} --api-base http://localhost:8000/v1",
                    file=sys.stderr,
                )
                return 1

            # Replace the entire pool with the single model
            config.llm.models = [LLMModelConfig(name=model_name)]
            config.llm.evaluator_models = config.llm.models.copy()
            config.llm.guide_models = config.llm.models.copy()

            # Force-propagate --api-base to all models.
            if args.api_base:
                config.llm.update_model_params({"api_base": args.api_base}, overwrite=True)

        else:
            if args.api_base:
                config.llm.api_base = args.api_base

        if args.model:
            print("Active models:")
            for i, model in enumerate(config.llm.models):
                print(f"  Model {i+1}: {model.name} (weight: {model.weight})")

        if args.api_base:
            print(f"Using API base: {config.llm.api_base}")

        if args.model or args.api_base:
            api_key = _resolve_api_key_from_env(parsed_env_vars)
            if api_key:
                config.llm.api_key = api_key
                config.llm.update_model_params({"api_key": api_key}, overwrite=True)

            # Ensure model overrides inherit shared LLM request settings.
            # Without this, overridden models can end up with retries/timeout=None.
            config.llm.update_model_params(
                {
                    "api_base": config.llm.api_base,
                    "api_key": config.llm.api_key,
                    "temperature": config.llm.temperature,
                    "top_p": config.llm.top_p,
                    "max_tokens": config.llm.max_tokens,
                    "timeout": config.llm.timeout,
                    "retries": config.llm.retries,
                    "retry_delay": config.llm.retry_delay,
                    "random_seed": config.llm.random_seed,
                    "reasoning_effort": config.llm.reasoning_effort,
                },
                overwrite=True,
            )

        # --codebase enables agentic mode
        if args.codebase:
            codebase_path = os.path.abspath(args.codebase)
            if not os.path.isdir(codebase_path):
                print(f"Error: Codebase directory '{args.codebase}' not found")
                return 1
            config.agentic.enabled = True
            config.agentic.codebase_root = codebase_path
            print(f"Agentic mode enabled (codebase: {codebase_path})")

        # --search sets the search algorithm
        if args.search:
            if not hasattr(config, "search"):
                from skydiscover.config import SearchConfig
                config.search = SearchConfig()
            config.search.type = args.search
            # Re-create database config if the search type changed, so the
            # correct subclass (e.g. EvoxDatabaseConfig) gets used.
            from skydiscover.config import _DB_CONFIG_BY_TYPE
            new_db_cls = _DB_CONFIG_BY_TYPE.get(args.search)
            if new_db_cls and not isinstance(config.search.database, new_db_cls):
                config.search.database = new_db_cls()
            print(f"Using search algorithm: {args.search}")

    # Initialize and run solution discovery
    try:
        # Handle external backends (openevolve, shinkaevolve, gepa)
        search_type = config.search.type if config and hasattr(config, "search") else None
        if search_type:
            from skydiscover.extras.external import is_external, get_runner, KNOWN_EXTERNAL, get_package_name
            if is_external(search_type):
                from skydiscover.config import build_output_dir
                output_dir = args.output or build_output_dir(search_type, args.initial_program)
                os.makedirs(output_dir, exist_ok=True)

                # Start monitor if enabled
                monitor_server = None
                monitor_callback = None
                feedback_reader = None
                if config.monitor.enabled:
                    try:
                        from skydiscover.extras.monitor import MonitorServer
                        from skydiscover.extras.monitor.callback import create_external_callback

                        monitor_server = MonitorServer(
                            host=config.monitor.host,
                            port=config.monitor.port,
                            max_solution_length=config.monitor.max_solution_length,
                        )
                        monitor_server.start()
                        monitor_callback = create_external_callback(monitor_server, time.time())

                        if config.monitor.summary_model:
                            monitor_server.configure_summary(
                                model=config.monitor.summary_model,
                                api_key=config.monitor.summary_api_key or "",
                                api_base=config.monitor.summary_api_base,
                                top_k=config.monitor.summary_top_k,
                                interval=config.monitor.summary_interval,
                            )
                        # Auto-enable human feedback when monitor is live
                        try:
                            from skydiscover.prompt.human_feedback import HumanFeedbackReader
                            feedback_path = getattr(config, "human_feedback_file", None) or os.path.join(
                                output_dir, "human_feedback.md"
                            )
                            feedback_mode = getattr(config, "human_feedback_mode", "append")
                            feedback_reader = HumanFeedbackReader(feedback_path, mode=feedback_mode)
                            monitor_server.set_feedback_reader(feedback_reader)
                            logger.info(f"Human feedback enabled. Feedback file: {feedback_path}")
                        except Exception as e:
                            logger.warning(f"Failed to set up human feedback: {e}")

                        url = f"http://localhost:{monitor_server.port}/"
                        print(f"\n{'─' * 50}\n  Live monitor: {url}\n{'─' * 50}\n", flush=True)
                        logger.info(f"Live monitor: {url}")
                    except Exception as e:
                        logger.warning(f"Failed to start monitor: {e}")

                try:
                    result = await get_runner(search_type)(
                        program_path=args.initial_program,
                        evaluator_path=args.evaluation_file,
                        config_obj=config,
                        iterations=args.iterations or config.max_iterations,
                        output_dir=output_dir,
                        monitor_callback=monitor_callback,
                        feedback_reader=feedback_reader,
                    )
                except ModuleNotFoundError as e:
                    pkg = get_package_name(search_type)
                    print(f"Error: {e}")
                    print(f"\nThe '{search_type}' backend requires its package to be installed.")
                    print(f"Install it with:  pip install {pkg}")
                    return 1
                finally:
                    if monitor_server:
                        try:
                            monitor_server.push_event({"type": "discovery_complete"})
                            time.sleep(0.5)
                            monitor_server.stop()
                        except Exception:
                            pass

                print(f"\nDiscovery process complete! Best score: {result.best_score:.4f}")
                return 0
            if search_type in KNOWN_EXTERNAL:
                pkg = get_package_name(search_type)
                print(f"Error: Search type '{search_type}' requires the '{pkg}' package. "
                      f"Install it with: pip install {pkg}")
                return 1

        main_controller = MainController(
            initial_program_path=args.initial_program,
            evaluation_file=args.evaluation_file,
            config=config,
            config_path=args.config if config is None else None,
            output_dir=args.output,
        )

        # Load from checkpoint if specified
        if args.checkpoint:
            if not os.path.exists(args.checkpoint):
                print(f"Error: Checkpoint directory '{args.checkpoint}' not found")
                return 1
            print(f"Loading checkpoint from {args.checkpoint}")

            main_controller.database.load(args.checkpoint)
            print(f"Checkpoint loaded successfully (iteration {main_controller.database.last_iteration})")

        # Run SkyDiscover
        best_program = await main_controller.run(
            iterations=args.iterations,
            checkpoint_path=args.checkpoint,
        )

        # Get the checkpoint path
        checkpoint_dir = os.path.join(main_controller.output_dir, "checkpoints")
        latest_checkpoint = None
        if os.path.exists(checkpoint_dir):
            checkpoints = [
                os.path.join(checkpoint_dir, d)
                for d in os.listdir(checkpoint_dir)
                if os.path.isdir(os.path.join(checkpoint_dir, d))
            ]
            if checkpoints:
                latest_checkpoint = sorted(
                    checkpoints, key=lambda x: int(x.split("_")[-1]) if "_" in x else 0
                )[-1]

        print("\nDiscovery process complete!")
        if best_program is None:
            print("No valid programs were found during discovery process.")
        else:
            print("Best program metrics:")
            for name, value in best_program.metrics.items():
                # Handle mixed types: format numbers as floats, others as strings
                if isinstance(value, (int, float)):
                    print(f"  {name}: {value:.4f}")
                else:
                    print(f"  {name}: {value}")

        if latest_checkpoint:
            print(f"\nLatest checkpoint saved at: {latest_checkpoint}")
            print(f"To resume, use: --checkpoint {latest_checkpoint}")

        return 0

    except Exception as e:
        print(f"Error: {str(e)}")
        import traceback

        traceback.print_exc()
        return 1


def main() -> int:
    """
    Main entry point

    Returns:
        Exit code
    """
    return asyncio.run(main_async())


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