import argparse
import sys
import tomllib
import pytest
from typing import Dict, List, Any

def load_manifest(path: str) -> Dict[str, Any]:
    """
    Loads a manifest file (expected to be in TOML format) from the given path.

    Args:
        path: The path to the TOML manifest file.

    Returns:
        A dictionary representing the loaded manifest.

    Raises:
        FileNotFoundError: If the manifest file does not exist.
        tomllib.TOMLDecodeError: If the manifest file is not valid TOML.
    """
    try:
        with open(path, 'rb') as f:
            return tomllib.load(f)
    except FileNotFoundError:
        print(f"Error: Manifest file not found at {path}", file=sys.stderr)
        raise
    except tomllib.TOMLDecodeError:
        print(f"Error: Could not decode TOML from {path}", file=sys.stderr)
        raise

def get_test_files(manifest: Dict[str, Any], category: str) -> List[str]:
    """
    Determines the list of test files based on the manifest and a specified category.

    Args:
        manifest: The loaded manifest dictionary.
        category: The category of tests to retrieve.

    Returns:
        A list of file paths corresponding to the tests in the given category.
        Returns an empty list if the category is not found or has no tests.
    """
    print(f"DEBUG: Looking for category '{category}' in manifest.", file=sys.stderr)
    files = manifest.get("categories", {}).get(category, {}).get("files", [])
    print(f"DEBUG: Found test files for category '{category}': {files}", file=sys.stderr)
    return files

def main():
    parser = argparse.ArgumentParser(
        description="Run tests with optional manifest and category filtering, passing additional pytest arguments.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
Example usage:
  python run_tests.py --manifest tests.toml --category unit -- --verbose --cov=my_module
  python run_tests.py --manifest tests.toml --category integration
  python run_tests.py --manifest tests.toml --category core
  python run_tests.py --manifest tests.toml # Runs tests from default_categories
  python run_tests.py -- --capture=no # Runs all tests with pytest args
"""
    )
    parser.add_argument(
        "--manifest",
        type=str,
        help="Path to the TOML manifest file containing test configurations."
    )
    parser.add_argument(
        "--category",
        type=str,
        help="Category of tests to run (e.g., 'unit', 'integration')."
    )

    # Parse known arguments for the script itself, then parse remaining args for pytest
    args, remaining_pytest_args = parser.parse_known_args(sys.argv[1:])

    selected_test_files = []
    manifest_data = None

    if args.manifest:
        try:
            manifest_data = load_manifest(args.manifest)
        except (FileNotFoundError, tomllib.TOMLDecodeError):
            # Error message already printed by load_manifest
            sys.exit(1)

        if args.category:
            # Case 1: --manifest and --category provided
            files = get_test_files(manifest_data, args.category)
            selected_test_files.extend(files)
        else:
            # Case 2: --manifest provided, but no --category
            # Load default categories from manifest['execution']['default_categories']
            default_categories = manifest_data.get("execution", {}).get("default_categories", [])
            if not default_categories:
                print(f"Error: --manifest provided without --category, and no 'default_categories' found in manifest '{args.manifest}'.", file=sys.stderr)
                parser.print_help(sys.stderr)
                sys.exit(1)

            print(f"DEBUG: Using default categories from manifest '{args.manifest}': {default_categories}", file=sys.stderr)
            for cat in default_categories:
                files = get_test_files(manifest_data, cat)
                selected_test_files.extend(files)
    elif args.category:
        # Case 3: --category provided without --manifest
        print("Error: --category requires --manifest to be specified.", file=sys.stderr)
        parser.print_help(sys.stderr)
        sys.exit(1)

    # Combine selected test files with any remaining pytest arguments that were not parsed by this script.
    # We also filter out the literal '--' if it was passed by the user to avoid pytest errors if it appears multiple times.
    pytest_command_args = selected_test_files + [arg for arg in remaining_pytest_args if arg != '--']

    # Filter out any empty strings that might have been included.
    final_pytest_args = [arg for arg in pytest_command_args if arg]

    # If no specific tests were selected from manifest/category and no manifest was provided,
    # and no other pytest args were given, pytest.main([]) runs default test discovery.
    print(f"Running pytest with arguments: {final_pytest_args}", file=sys.stderr)
    sys.exit(pytest.main(final_pytest_args))

if __name__ == "__main__":
    main()
