import argparse
import re
from collections import defaultdict
from typing import List, Tuple, DefaultDict, Set, Optional


def parse_requirement(line: str) -> Optional[Tuple[str, Optional[str]]]:
    """
    Parses a single line from a requirements.txt file.

    Args:
        line: The line to parse.

    Returns:
        A tuple containing the package name and version, or (None, None) if the line is empty or a comment.
    """
    line = line.strip()
    if not line or line.startswith('#'):
        return None, None

    # This regex is designed to parse a requirement line into a package name and an optional version specifier.
    # It is not fully compliant with PEP 508, but it is sufficient for most common cases.
    match = re.match(r'^([a-zA-Z0-9_.-]+)((?:[<>=!~]=?|==)?[a-zA-Z0-9_.-]+)?$', line)
    if match:
        package_name = match.group(1)
        version_specifier = match.group(2) if match.group(2) else None
        return package_name, version_specifier

    return line, None


def read_requirements_files(file_paths: List[str]) -> DefaultDict[str, Set[Optional[str]]]:
    """
    Reads multiple requirements.txt files and merges them into a single dictionary.

    Args:
        file_paths: A list of paths to the requirements.txt files.

    Returns:
        A dictionary where the keys are package names and the values are sets of version specifiers.
    """
    merged_requirements: DefaultDict[str, Set[Optional[str]]] = defaultdict(set)

    for file_path in file_paths:
        with open(file_path, 'r') as f:
            for line in f:
                package_name, version_specifier = parse_requirement(line)
                if package_name:
                    merged_requirements[package_name].add(version_specifier if version_specifier else 'unspecified')

    return merged_requirements


def write_merged_requirements(output_path: str, merged_requirements: DefaultDict[str, Set[Optional[str]]],
                              version: str = None) -> None:
    """
    Writes the merged requirements to a file, flagging conflicts.

    Args:
        output_path: The path to the output file.
        merged_requirements: A dictionary of merged requirements.
        version: If there is a conflict, use this version.
    """
    with open(output_path, 'w') as f:
        for package_name in sorted(merged_requirements.keys()):
            version_specifiers = merged_requirements[package_name]

            if len(version_specifiers) > 1:
                # If there are multiple version specifiers for a package, it's a conflict.
                # We comment out all the conflicting versions and add a note to the user.

                version_strings = sorted([v for v in version_specifiers if v])
                if not version:
                    if None in version_specifiers:
                        version_strings.append('unspecified')

                    f.write(
                        f'# CONFLICT: Found multiple versions for {package_name}: {" and ".join(version_strings)}\n')
                    for version_specifier in sorted(version_specifiers, key=lambda v: v or ''):
                        f.write(f'# {package_name}{version_specifier if version_specifier != 'unspecified' else ""}\n')
                # Prefer the latest version if specified.
                if version:
                    f.write(
                        f'# CONFLICT: Found multiple versions for {package_name}: {" and ".join(version_strings)}. Keeping {version} version\n'
                    )
                    if 'unspecified' in version_strings and version == 'latest':
                        f.write(
                            f'{package_name}\n'
                        )
                    else:
                        f.write(
                            f'{package_name}{version_strings[(len(version_strings) - 1) if version == 'latest' else 0]}\n'
                        )

            else:
                # If there is only one version specifier, we can write it to the file.
                version_specifier = list(version_specifiers)[0]
                f.write(f'{package_name}{version_specifier if version_specifier != 'unspecified' else ''}\n')


def merge_requirements(file_paths: List[str], output_path: str, version: str = None) -> None:
    """
    Merges multiple requirements.txt files into a single file,
    flagging conflicts and sorting the requirements alphabetically.
    """
    merged_requirements = read_requirements_files(file_paths)
    write_merged_requirements(output_path, merged_requirements, version)


def main():
    """The main entry point of the script."""
    parser = argparse.ArgumentParser(description='Merge multiple requirements.txt files.')
    parser.add_argument('files', nargs='+', help='List of requirements.txt files to merge.')
    parser.add_argument('-o', '--output', default='merged_requirements.txt', help='Output file name.')
    parser.add_argument(
        '-s',
        '--strategy',
        choices=['latest', 'oldest'],
        default=None,
        help='Resolve conflicts using latest or oldest version'
    )
    args = parser.parse_args()

    merge_requirements(args.files, args.output, args.strategy)
    print(f"Successfully merged requirements into {args.output}")


if __name__ == '__main__':
    main()
