"""Export locked production requirements for the Kylin ARM64 CPU target."""

import argparse
from pathlib import Path
import subprocess

from packaging.markers import default_environment
from packaging.requirements import Requirement


ROOT = Path(__file__).resolve().parents[1]


def target_environment():
    environment = default_environment()
    environment.update({
        'implementation_name': 'cpython',
        'os_name': 'posix',
        'platform_machine': 'aarch64',
        'platform_python_implementation': 'CPython',
        'platform_system': 'Linux',
        'python_full_version': '3.14.5',
        'python_version': '3.14',
        'sys_platform': 'linux',
    })
    return environment


def filter_export(exported):
    environment = target_environment()
    selected = []
    cpu_torch = False
    for line in exported.splitlines():
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        requirement = Requirement(line)
        if requirement.marker and not requirement.marker.evaluate(environment):
            continue
        requirement.marker = None
        if requirement.name == 'torch' and '+cpu' in str(requirement.specifier):
            cpu_torch = True
        selected.append(str(requirement))
    if not cpu_torch:
        raise RuntimeError('target_cpu_torch_missing')
    if any(line.startswith(('cuda-', 'nvidia-', 'triton==')) for line in selected):
        raise RuntimeError('unexpected_cuda_dependency')
    return '\n'.join(selected) + '\n'


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('output', type=Path)
    args = parser.parse_args()
    export = subprocess.run(
        ['uv', '--cache-dir', str(ROOT / '.uv-cache'), 'export', '--locked',
         '--offline', '--no-dev', '--no-emit-project', '--no-hashes',
         '--no-annotate', '--format', 'requirements-txt'],
        cwd=ROOT, check=True, capture_output=True, text=True,
    )
    requirements = filter_export(export.stdout)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(requirements, encoding='utf-8')
    print(f'Wrote {len(requirements.splitlines())} locked ARM64 CPU requirements.')


if __name__ == '__main__':
    main()