import shutil
import tensorrt_llm
from pathlib import Path

trtllm_path = Path(tensorrt_llm.__file__).parent
target_dir = trtllm_path / "models"

print(f"TensorRT-LLM path: {trtllm_path}")
print(f"Target models directory: {target_dir}")

target_dir.mkdir(parents=True, exist_ok=True)

patch_dir = Path("./patch")

patch_files = list(patch_dir.glob('*'))
if patch_files:
    print(f"Copying {len(patch_files)} patch file(s) to tensorrt_llm/models")

    for patch_file in patch_files:
        target_path = target_dir / patch_file.name
        if patch_file.is_file():
            shutil.copy2(patch_file, target_path)
            print(f"  Copied: {patch_file.name}")
        elif patch_file.is_dir():
            if target_path.exists():
                shutil.rmtree(target_path)
            shutil.copytree(patch_file, target_path)
            print(f"  Copied directory: {patch_file.name}")

    print(f"✓ Patch files copied successfully")
else:
    print(f"⚠ No patch files found in {patch_dir}")