"""
``@ctrl.main()`` decorator for Investigation experiment entry points.

Wraps a ``main()`` function so it can operate in two modes without any changes
to user code:

* **Normal mode** — standard CLI invocation, arguments parsed by
  `tyro <https://brentyi.github.io/tyro/>`_ as usual.
* **Investigation mode** — triggered when ``--doe-config`` and ``--commit`` are
  present on the command line; the decorator processes the DoE JSON, generates
  case databases and Slurm shell scripts, then exits without calling the user
  function.
"""
import sys
import json
from pathlib import Path
from typing import Callable, Optional
from functools import wraps


def main(doe_config: Optional[str] = None, commit: Optional[str] = None):
    """Decorator factory for a Investigation experiment entry point.

    Wraps a ``main()`` function so it operates in two modes:

    **Normal mode** — standard CLI, arguments parsed by tyro as usual:

    .. code-block:: bash

        poetry run python -m src.entry --seed 42 --name MyExperiment

    **Investigation mode** — triggered when ``--doe-config`` and ``--commit``
    are present; processes the DoE JSON and generates shell scripts, then
    exits without calling the user function:

    .. code-block:: bash

        poetry run python -m src.entry --doe-config doe.json --commit abc123

    Args:
        doe_config: Path to DoE JSON file.  If ``None``, the decorator reads
            ``--doe-config`` from ``sys.argv`` and removes it before passing
            control to the user function.
        commit: Git commit hash.  If ``None``, read from ``--commit`` in
            ``sys.argv``.

    Returns:
        A decorator that wraps the user's ``main()`` function.

    Example:
        .. code-block:: python

            import investigation as ctrl
            import tyro
            from src.args import Args

            @ctrl.main()
            def main() -> None:
                args = tyro.cli(Args)
                from src.train import train
                train(args)

            if __name__ == '__main__':
                main()
    """

    def decorator(func: Callable) -> Callable:

        @wraps(func)
        def wrapper(*args, **kwargs):
            # Parse command line arguments to check for --doe-config and --commit
            parsed_doe_config = doe_config
            parsed_commit = commit

            # Import parse_unknown_args early (no heavy dependencies)
            from investigation.cli.utils import parse_unknown_args

            # Check if --doe-config and --commit are in sys.argv
            argv = sys.argv[1:]
            unknown_args = parse_unknown_args(argv)

            if 'doe_config' in unknown_args and parsed_doe_config is None:
                parsed_doe_config = unknown_args['doe_config']
                # Remove from sys.argv so it doesn't interfere with user's arg parsing
                if '--doe-config' in sys.argv:
                    idx = sys.argv.index('--doe-config')
                    sys.argv.pop(idx)  # Remove --doe-config
                    # Check if there's a value after --doe-config
                    if idx < len(sys.argv) and not sys.argv[idx].startswith('--'):
                        sys.argv.pop(idx)  # Remove the value

            if 'commit' in unknown_args and parsed_commit is None:
                parsed_commit = unknown_args['commit']
                # Remove from sys.argv so it doesn't interfere with user's arg parsing
                if '--commit' in sys.argv:
                    idx = sys.argv.index('--commit')
                    sys.argv.pop(idx)  # Remove --commit
                    # Check if there's a value after --commit
                    if idx < len(sys.argv) and not sys.argv[idx].startswith('--'):
                        sys.argv.pop(idx)  # Remove the value

            # If both doe_config and commit are provided, process investigation workflow
            if parsed_doe_config and parsed_commit:
                # Import heavy dependencies only when needed
                from investigation import storage_path, cwd
                from investigation.command import commit_entry

                doe_config_path = Path(parsed_doe_config)
                if not doe_config_path.exists():
                    print(f"Error: DoE configuration file {doe_config_path} not found.")
                    sys.exit(1)

                # Read DoE configuration
                doe_config_data = json.loads(doe_config_path.read_text())

                # Commit the configuration
                print(f"Processing DoE configuration: {doe_config_path}")
                print(f"Commit hash: {parsed_commit}")

                # Create a temporary file with the DoE config for commit_entry
                temp_config_name = doe_config_path.stem
                temp_config_path = cwd / f"{temp_config_name}.json"

                # If temp config doesn't exist, create it
                if not temp_config_path.exists():
                    temp_config_path.write_text(json.dumps(doe_config_data, indent=4))

                # Commit the entry
                try:
                    commit_entry(temp_config_name, parsed_commit)
                    print(f"DoE configuration committed successfully.")
                    print(f"Cases and shell scripts generated in {storage_path / parsed_commit}")
                except Exception as e:
                    print(f"Error during commit: {e}")
                    sys.exit(1)
                finally:
                    # Clean up temp file after commit (success or failure)
                    if temp_config_path.exists() and temp_config_path.stem == temp_config_name:
                        temp_config_path.unlink()

                # Don't call the user's function when in investigation mode
            else:
                # Otherwise, call the user's function normally
                return func(*args, **kwargs)

        return wrapper

    return decorator
