"""Command-line interface for codeflash-python.

Thin wiring layer: arg parsing, config loading, and delegation to
:class:`PythonOptimizer` / :class:`PythonFunctionOptimizer`.  No
business logic lives here.
"""

from __future__ import annotations

import logging
import sys
from argparse import SUPPRESS, ArgumentParser, Namespace
from functools import lru_cache
from pathlib import Path
from typing import TYPE_CHECKING, Any, NoReturn

if TYPE_CHECKING:
    from .._model import FunctionToOptimize
    from ..benchmarking.models import BenchmarkKey

log = logging.getLogger(__name__)


def parse_args(
    argv: list[str] | None = None,
) -> Namespace:
    """Parse CLI arguments.

    Returns a validated :class:`~argparse.Namespace`.
    """
    parser = _build_parser()
    args, unknown = parser.parse_known_args(argv)
    # Let pytest see leftover args when invoked via subprocess.
    sys.argv[:] = [sys.argv[0], *unknown]
    return args


@lru_cache(maxsize=1)
def _build_parser() -> ArgumentParser:
    parser = ArgumentParser(
        prog="codeflash",
        description="Optimize Python functions with AI.",
    )

    # -- Target selection -----------------------------------------------
    parser.add_argument(
        "--file",
        type=str,
        help="Optimize only this file.",
    )
    parser.add_argument(
        "--function",
        type=str,
        help=(
            "Optimize only this function within the given file."
            " Requires --file."
        ),
    )
    parser.add_argument(
        "--all",
        nargs="?",
        const="",
        default=SUPPRESS,
        help=(
            "Optimize all functions.  Optionally pass a starting"
            " directory (defaults to module-root)."
        ),
    )

    # -- Paths & config -------------------------------------------------
    parser.add_argument(
        "--module-root",
        type=str,
        help="Root directory of the source code to optimize.",
    )
    parser.add_argument(
        "--tests-root",
        type=str,
        help="Root directory of the test suite.",
    )
    parser.add_argument(
        "--benchmarks-root",
        type=str,
        help="Root of pytest-benchmark tests.",
    )
    parser.add_argument(
        "--config-file",
        type=str,
        help="Path to pyproject.toml with [tool.codeflash] config.",
    )

    # -- Mode flags -----------------------------------------------------
    parser.add_argument(
        "--no-gen-tests",
        action="store_true",
        help="Only use existing tests (skip test generation).",
    )
    parser.add_argument(
        "--benchmark",
        action="store_true",
        help="Run benchmarks and rank by impact.",
    )
    parser.add_argument(
        "--no-pr",
        action="store_true",
        help="Don't open a PR; only update code locally.",
    )
    parser.add_argument(
        "--effort",
        type=str,
        choices=["low", "medium", "high"],
        default="medium",
        help="Optimization effort level.",
    )

    # -- AI service -----------------------------------------------------
    parser.add_argument(
        "--server",
        type=str,
        choices=["local", "prod"],
        help=(
            "AI service: 'local' for localhost:8000,"
            " 'prod' for app.codeflash.ai."
        ),
    )

    # -- Misc -----------------------------------------------------------
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Enable debug logging.",
    )
    parser.add_argument(
        "--subagent",
        action="store_true",
        help="Subagent mode: non-interactive defaults.",
    )
    parser.add_argument(
        "--pytest-cmd",
        type=str,
        help="Pytest executable (default: pytest).",
    )

    return parser


def load_config(args: Namespace) -> Namespace:
    """Merge pyproject.toml ``[tool.codeflash]`` config into *args*.

    CLI flags take precedence over config-file values.  Paths are
    resolved relative to the config file's parent directory.
    """
    import os  # noqa: PLC0415

    from ._config import parse_config_file  # noqa: PLC0415

    if args.server:
        os.environ["CODEFLASH_AIS_SERVER"] = args.server

    if args.subagent:
        os.environ["CODEFLASH_SUBAGENT_MODE"] = "true"
        args.no_pr = True
        args.effort = "low"

    cfg, config_path = parse_config_file(
        Path(args.config_file) if args.config_file else None,
    )

    _merge_config(args, cfg)
    _resolve_paths(args, config_path)
    _validate(args)

    return args


def _merge_config(args: Namespace, cfg: dict[str, Any]) -> None:
    """Apply config-file values where CLI didn't set them."""
    config_keys = [
        "module_root",
        "tests_root",
        "benchmarks_root",
        "ignore_paths",
        "pytest_cmd",
        "formatter_cmds",
        "disable_telemetry",
        "git_remote",
    ]
    for key in config_keys:
        cli_val = getattr(args, key, None)
        if cli_val is None and key in cfg:
            setattr(args, key, cfg[key])


def _resolve_paths(args: Namespace, config_path: Path) -> None:
    """Resolve module_root, tests_root, project_root to absolute."""
    _resolve_root_paths(args)
    _resolve_project_root(args, config_path)
    _resolve_target_paths(args)


def _resolve_root_paths(args: Namespace) -> None:
    """Resolve module_root, tests_root, benchmarks_root."""
    if args.module_root is not None:
        args.module_root = Path(args.module_root).resolve()
    if hasattr(args, "tests_root") and args.tests_root is not None:
        args.tests_root = Path(args.tests_root).resolve()
    if hasattr(args, "benchmarks_root") and args.benchmarks_root is not None:
        args.benchmarks_root = Path(args.benchmarks_root).resolve()


def _resolve_project_root(
    args: Namespace,
    config_path: Path,
) -> None:
    """Derive project_root and test_project_root."""
    # The project root is the directory containing pyproject.toml.
    # This correctly handles nested layouts like futurehouse where
    # module_root (src/aviary) is several levels below the config.
    args.project_root = config_path.parent.resolve()
    args.test_project_root = args.project_root


def _resolve_target_paths(args: Namespace) -> None:
    """Resolve --all, --file, and ensure ignore_paths exists."""
    if hasattr(args, "all"):
        if args.all == "":
            args.all = args.module_root
        elif args.all is not None:
            args.all = Path(args.all).resolve()
    if args.file:
        args.file = Path(args.file).resolve()
    if not hasattr(args, "ignore_paths") or args.ignore_paths is None:
        args.ignore_paths = []


def _validate(args: Namespace) -> None:
    """Validate required args and flag combinations."""
    if args.module_root is None:
        _exit("--module-root is required (set in pyproject.toml or CLI)")
    if not Path(args.module_root).is_dir():
        _exit(f"--module-root {args.module_root} is not a directory")
    if args.function and not args.file:
        _exit("--function requires --file")
    if args.file and not Path(args.file).exists():
        _exit(f"--file {args.file} does not exist")
    if not hasattr(args, "all") and not args.file:
        # Default to --all when neither --file nor --all is given.
        args.all = args.module_root


def _discover_target_functions(
    args: Namespace,
) -> dict[Path, list[Any]]:
    """Discover functions to optimize based on CLI target flags."""
    from ..analysis._discovery import (  # noqa: PLC0415
        find_all_functions_in_file,
        get_all_files_and_functions,
    )

    if args.file:
        file_path = Path(args.file)
        file_to_funcs = find_all_functions_in_file(file_path)
        if args.function:
            file_to_funcs = _filter_to_function(
                file_to_funcs,
                file_path,
                args.function,
            )
        return file_to_funcs

    # --all mode.
    target = getattr(args, "all", args.module_root)
    return get_all_files_and_functions(
        Path(target),
        args.ignore_paths,
    )


def _filter_to_function(
    file_to_funcs: dict[Path, list[Any]],
    file_path: Path,
    function_name: str,
) -> dict[Path, list[Any]]:
    """Filter to a single function (supports Class.method syntax)."""
    parts = function_name.split(".")
    if len(parts) == 2:  # noqa: PLR2004
        class_name, func_name = parts
    else:
        class_name, func_name = None, parts[0]

    for fn in file_to_funcs.get(file_path, []):
        if fn.function_name == func_name and (
            class_name is None or fn.top_level_parent_name == class_name
        ):
            return {file_path: [fn]}

    _exit(
        f"Function {function_name!r} not found in {file_path}",
    )
    return {}  # unreachable, but satisfies type checker


def _collect_benchmarks(
    args: Namespace,
    file_to_funcs: dict[Path, list[FunctionToOptimize]],
    tests_root: str,
) -> tuple[
    dict[str, dict[BenchmarkKey, int]] | None,
    dict[BenchmarkKey, int] | None,
    Path | None,
]:
    """Run benchmark collection if ``--benchmark`` is enabled."""
    benchmarks_root = getattr(args, "benchmarks_root", None)
    if not getattr(args, "benchmark", False) or not benchmarks_root:
        return None, None, None

    from ._orchestrator import run_benchmarks  # noqa: PLC0415

    return run_benchmarks(
        file_to_funcs,
        Path(benchmarks_root),
        Path(tests_root),
        args.project_root,
    )


def main(argv: list[str] | None = None) -> int:  # noqa: C901, PLR0912, PLR0915
    """CLI entry point for ``codeflash`` / ``python -m codeflash_python``."""
    from codeflash_core import (  # noqa: PLC0415
        AIClient,
        AIServiceConnectionError,
        AIServiceError,
        InvalidAPIKeyError,
        PlatformClient,
        init_telemetry,
    )
    from codeflash_core import (  # noqa: PLC0415
        __version__ as codeflash_core_version,
    )

    from ..test_discovery.discovery import (  # noqa: PLC0415
        discover_unit_tests,
    )
    from ..testing.models import TestConfig  # noqa: PLC0415
    from ._context import OptimizationContext  # noqa: PLC0415
    from ._function_optimizer import (  # noqa: PLC0415
        PythonFunctionOptimizer,
    )
    from ._optimizer import PythonOptimizer  # noqa: PLC0415
    from ._plugin import PythonPlugin  # noqa: PLC0415

    # 1. Parse args and load config.
    args = parse_args(argv)

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)

    try:
        args = load_config(args)
    except (ValueError, TypeError) as e:
        _exit(str(e))
    log.info("Project root: %s", args.project_root)
    log.info("Module root:  %s", args.module_root)

    # 2. Discover functions.
    file_to_funcs = _discover_target_functions(args)
    total = sum(len(fns) for fns in file_to_funcs.values())
    if total == 0:
        log.info("No optimizable functions found")
        return 0
    log.info(
        "Found %d function(s) across %d file(s)", total, len(file_to_funcs)
    )

    # 3. Set up test config and discover tests.
    tests_root = getattr(args, "tests_root", None)
    if tests_root is None:
        _exit("--tests-root is required (set in pyproject.toml or CLI)")

    test_cfg = TestConfig(
        tests_root=Path(tests_root),
        tests_project_rootdir=args.test_project_root,
        project_root_path=args.project_root,
        test_framework="pytest",
        pytest_cmd=getattr(args, "pytest_cmd", None) or "pytest",
        module_root=args.module_root,
    )

    # Clean up leftover instrumented files/dirs from previous
    # (possibly crashed) runs *before* test discovery, so that
    # stale temp dirs don't confuse pytest collection.
    from ._orchestrator import (  # noqa: PLC0415
        cleanup_paths,
        find_leftover_instrumented_test_files,
    )

    leftover = find_leftover_instrumented_test_files(Path(tests_root))
    if leftover:
        log.debug(
            "Removing %d leftover instrumented file(s)/dir(s)",
            len(leftover),
        )
        cleanup_paths(leftover)

    import time as _time  # noqa: PLC0415

    _t0 = _time.time()
    function_to_tests, n_tests, n_replay = discover_unit_tests(
        test_cfg,
        file_to_funcs_to_optimize=file_to_funcs,
    )
    _elapsed = _time.time() - _t0
    n_unit = n_tests - n_replay
    log.info(
        "Discovered %d existing unit test%s and %d replay test%s"
        " in %.1fs at %s",
        n_unit,
        "" if n_unit == 1 else "s",
        n_replay,
        "" if n_replay == 1 else "s",
        _elapsed,
        tests_root,
    )

    # 4. Wire up optimizers.
    from .._configuration import PythonConfiguration  # noqa: PLC0415
    from .._state import PythonState  # noqa: PLC0415

    cfg = PythonConfiguration(
        project_root=args.project_root,
        tests_root=Path(tests_root),
        test_framework="pytest",
        pytest_cmd=getattr(args, "pytest_cmd", None) or "pytest",
        module_root=args.module_root,
        no_pr=getattr(args, "no_pr", False),
        no_gen_tests=getattr(args, "no_gen_tests", False),
        git_remote=getattr(args, "git_remote", "origin"),
    )
    state = PythonState(cfg=cfg)
    plugin = PythonPlugin(configuration=cfg, state=state)

    with PlatformClient() as platform_client, AIClient() as ai_client:
        # 4a. Validate the API key (fail fast on invalid keys).
        user_id: str | None = None
        try:
            user_id = platform_client.validate_api_key()
        except InvalidAPIKeyError:
            log.error(  # noqa: TRY400
                "Invalid API key."
                " Set CODEFLASH_API_KEY or generate one at"
                " https://app.codeflash.ai/app/apikeys",
            )
            return 1
        except (AIServiceError, AIServiceConnectionError):
            log.warning(
                "Could not validate API key (network error);"
                " proceeding anyway",
            )

        # 4b. Initialize telemetry (PostHog + Sentry).
        disable_telemetry = getattr(
            args,
            "disable_telemetry",
            False,
        )
        init_telemetry(
            platform_client,
            version=codeflash_core_version,
            enabled=not disable_telemetry,
            user_id=user_id,
        )

        # 4c. Run benchmarks if requested.
        (
            function_benchmark_timings,
            total_benchmark_timings,
            replay_tests_dir,
        ) = _collect_benchmarks(args, file_to_funcs, tests_root)

        opt_ctx = OptimizationContext(
            plugin=plugin,
            project_root=args.project_root,
            test_cfg=test_cfg,
            ai_client=ai_client,
        )
        fn_optimizer = PythonFunctionOptimizer(
            ctx=opt_ctx,
            function_to_tests=function_to_tests,
            replay_tests_dir=replay_tests_dir,
            no_gen_tests=args.no_gen_tests,
        )

        # PR creation setup.
        import git as _git  # noqa: PLC0415

        git_repo: _git.Repo | None = None
        pr_platform_client: PlatformClient | None = None
        no_pr = getattr(args, "no_pr", False)
        git_remote = getattr(args, "git_remote", None) or "origin"
        if not no_pr:
            try:
                git_repo = _git.Repo(
                    str(args.project_root),
                    search_parent_directories=True,
                )
            except _git.InvalidGitRepositoryError:
                log.warning(
                    "Not a git repository; PR creation disabled",
                )
                no_pr = True
            else:
                pr_platform_client = platform_client

        project_optimizer = PythonOptimizer(
            plugin=plugin,
            project_root=args.project_root,
            no_pr=no_pr,
            git_remote=git_remote,
            platform_client=pr_platform_client,
            git_repo=git_repo,
        )

        # 5. Run the optimization loop.
        import asyncio  # noqa: PLC0415

        try:
            results = asyncio.run(
                project_optimizer.run(
                    file_to_funcs=file_to_funcs,
                    optimize_fn=fn_optimizer.optimize,
                    tests_root=Path(tests_root),
                    function_benchmark_timings=function_benchmark_timings,
                    total_benchmark_timings=total_benchmark_timings,
                ),
            )
        finally:
            _cleanup_run(replay_tests_dir, Path(tests_root))

    # 6. Report.
    succeeded = sum(1 for r in results if r.success)
    log.info(
        "%d of %d function(s) optimized",
        succeeded,
        len(results),
    )
    for r in results:
        status = "OK" if r.success else "SKIP"
        log.info(
            "  [%s] %s — %s",
            status,
            r.function.qualified_name,
            r.message,
        )

    return 0 if succeeded > 0 else 1


def _cleanup_run(
    replay_tests_dir: Path | None,
    tests_root: Path,
) -> None:
    """Remove run-level temp paths (replay dir, trace files)."""
    from ._orchestrator import cleanup_paths  # noqa: PLC0415

    paths: list[Path | None] = [replay_tests_dir]
    if tests_root.exists():
        paths.extend(tests_root.glob("*.trace"))
    cleanup_paths(paths)


def _exit(msg: str) -> NoReturn:
    """Print an error and exit."""
    log.error(msg)
    sys.exit(1)
