from __future__ import annotations

import difflib
import re
from pathlib import Path

from .ingest import iter_text_files


def generate_patch_diff(root: Path) -> str:
    diff_parts: list[str] = []
    for relative_path, original in iter_text_files(root):
        transformed = transform_text(relative_path, original)
        if transformed == original:
            continue
        diff = difflib.unified_diff(
            original.splitlines(keepends=True),
            transformed.splitlines(keepends=True),
            fromfile=f"a/{relative_path}",
            tofile=f"b/{relative_path}",
        )
        diff_parts.extend(diff)

    if not diff_parts:
        return "# No deterministic patch was generated. Review manual findings in the migration report.\n"
    return "".join(diff_parts)


def transform_text(relative_path: str, text: str) -> str:
    path = Path(relative_path)
    lower_name = path.name.lower()
    suffix = path.suffix.lower()

    if suffix == ".py":
        return _transform_python(text)
    if lower_name.startswith("dockerfile"):
        return _transform_dockerfile(text)
    if suffix in {".sh", ".bash", ".zsh", ".yaml", ".yml", ".txt", ".md"}:
        return _transform_shellish(text)
    return text


def _transform_python(text: str) -> str:
    changed = text
    needs_device = bool(
        re.search(r"\.cuda\s*\(\s*\)", changed)
        or re.search(r"\.to\(\s*['\"]cuda['\"]\s*\)", changed)
        or re.search(r"torch\.device\(\s*['\"]cuda['\"]\s*\)", changed)
    )
    if needs_device and "import torch" in changed and "_rocmport_device" not in changed:
        changed = _insert_device_helper(changed)

    changed = re.sub(r"\.cuda\s*\(\s*\)", ".to(_rocmport_device)", changed)
    changed = re.sub(r"\.to\(\s*['\"]cuda['\"]\s*\)", ".to(_rocmport_device)", changed)
    changed = re.sub(r"torch\.device\(\s*['\"]cuda['\"]\s*\)", "_rocmport_device", changed)
    return changed


def _insert_device_helper(text: str) -> str:
    lines = text.splitlines()
    insert_at = 0
    for index, line in enumerate(lines):
        if line.startswith("import ") or line.startswith("from "):
            insert_at = index + 1
    helper = [
        "",
        "# ROCm PyTorch exposes AMD GPUs through the torch.cuda namespace.",
        '_rocmport_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")',
    ]
    return "\n".join(lines[:insert_at] + helper + lines[insert_at:]) + ("\n" if text.endswith("\n") else "")


def _transform_dockerfile(text: str) -> str:
    changed = re.sub(
        r"(?im)^\s*FROM\s+nvidia/cuda:[^\n]+",
        "FROM vllm/vllm-openai-rocm:latest",
        text,
    )
    changed = _transform_shellish(changed)
    return changed


def _transform_shellish(text: str) -> str:
    changed = text.replace("nvidia-smi", "rocm-smi")
    changed = changed.replace("NVIDIA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES")
    changed = changed.replace("CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES")
    changed = changed.replace("NVIDIA_DRIVER_CAPABILITIES", "ROCM_VISIBLE_DEVICES")
    return changed