import json
import re
from pathlib import Path

class GeneratorCore:
    def __init__(self, config):
        self._cfg = config

    @classmethod
    def from_file(cls, file_path):
        with open(file_path, "r", encoding="utf-8") as f:
            return cls(json.load(f))

    def _sub(self, text, variables):
        def replace(match):
            key = match.group(1)
            if key not in variables:
                raise KeyError(f"Unknown template variable: {{{key}}}")
            return str(variables[key])
        return re.sub(r"\{(\w+)\}", replace, text)

    def _sub_all(self, arr, variables):
        return [self._sub(s, variables) for s in arr]

    def get_base_image(self, distro, variant, has_cuda):
        d = self._cfg["distros"].get(distro)
        if not d:
            raise ValueError(f"Unknown distro: {distro}")
        
        if has_cuda:
            return self._sub(self._cfg["cuda_base_image"], {
                "cuda_version": d["cuda_version"],
                "ubuntu_image_suffix": d["ubuntu_image_suffix"]
            })
            
        v = self._cfg["variants"].get(variant)
        if not v:
            raise ValueError(f"Unknown variant: {variant}")
            
        return self._sub(v["base_image"], {"distro": distro})

    def resolve_ros_package_apt(self, pkg_key, distro):
        pkg = self._cfg["ros_packages"].get(pkg_key)
        if not pkg or pkg.get("switches_base_image"):
            return []
            
        if "skip_on_distros" in pkg and distro in pkg["skip_on_distros"]:
            return []
            
        variables = {"distro": distro}
        names = self._sub_all(pkg.get("apt", []), variables)
        
        extra = pkg.get("extra_apt_on_distros")
        if extra and distro in extra.get("distros", []):
            names.extend(self._sub_all(extra.get("apt", []), variables))
            
        return names

    def build_dockerfile(self, config):
        distro = config["distro"]
        variant = config["variant"]
        packages = sorted(config["packages"])
        tools = sorted(config["tools"])
        username = config["username"]
        uid = config["uid"]
        workspace = config["workspace"]
        user_type = config.get("userType", config.get("user_type"))
        
        is_root = user_type == "root"
        has_cuda = "cuda" in packages or "tensorrt" in packages
        d = self._cfg["distros"][distro]
        variables = {
            "distro": distro,
            "cuda_version": d["cuda_version"],
            "ubuntu_image_suffix": d["ubuntu_image_suffix"]
        }
        
        L = []
        def ln(s): L.append(s)
        def gap(): L.append("")
        
        ln("# =============================================================")
        ln(f"# ROS2 {d['label']} Dockerfile")
        ln(f"# Generated by ros2-dockergen v{self._cfg['version']}")
        ln("# =============================================================")
        gap()
        ln(f"FROM {self.get_base_image(distro, variant, has_cuda)}")
        gap()
        
        if has_cuda:
            ln("# ── Install ROS2 on top of CUDA base ───────────────────────")
            ln("ENV DEBIAN_FRONTEND=noninteractive")
            ln("RUN apt-get update && apt-get install -y \\")
            for p in self._cfg["cuda_ros_install_apt"]:
                ln(f"    {p} \\")
            ln(f"    && curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key \\")
            ln("    -o /usr/share/keyrings/ros-archive-keyring.gpg && \\")
            ln("    echo \"deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] \\")
            ln("    http://packages.ros.org/ros2/ubuntu $(lsb_release -cs) main\" \\")
            ln("    > /etc/apt/sources.list.d/ros2.list && \\")
            ln(f"    apt-get update && apt-get install -y ros-{distro}-{variant} && \\")
            ln("    rm -rf /var/lib/apt/lists/*")
            gap()
            
        ln("ENV DEBIAN_FRONTEND=noninteractive")
        ln("SHELL [\"/bin/bash\", \"-c\"]")
        gap()
        
        if not has_cuda and variant == "desktop-full":
            extra = self._sub_all(self._cfg["variants"]["desktop-full"].get("extra_apt", []), {"distro": distro})
            ln("# ── Upgrade to desktop-full ─────────────────────────────────")
            ln("RUN apt-get update && apt-get install -y \\")
            for p in extra:
                ln(f"    {p} \\")
            ln("    && rm -rf /var/lib/apt/lists/*")
            gap()
            
        if "locale" in tools:
            locale = self._cfg["tools"]["locale"]
            ln("# ── Locale ───────────────────────────────────────────────────")
            ln("RUN apt-get update && apt-get install -y {} && \\".format(" ".join(locale["apt"])))
            ln(f"    {locale['post_install_run']} && \\")
            ln("    rm -rf /var/lib/apt/lists/*")
            for k, v in locale.get("env", {}).items():
                ln(f"ENV {k}={v}")
            gap()
            
        sys_apt = list(self._cfg["system_apt_always"])
        for tool_key in tools:
            if tool_key in ["locale", "zsh", "ssh"]:
                continue
            tool = self._cfg["tools"].get(tool_key)
            if tool and tool.get("apt") and not tool.get("ros_build_tool"):
                sys_apt.extend(tool["apt"])
        if not is_root:
            sys_apt.append("sudo")
            
        ln("# ── System packages ──────────────────────────────────────────")
        ln("RUN apt-get update && apt-get install -y \\")
        for p in sys_apt:
            ln(f"    {p} \\")
        ln("    && rm -rf /var/lib/apt/lists/*")
        gap()
        
        build_apt = list(self._cfg["ros_build_tools_always"])
        for tool_key in tools:
            tool = self._cfg["tools"].get(tool_key)
            if tool and tool.get("ros_build_tool") and tool.get("apt"):
                build_apt.extend(tool["apt"])
                
        ln("# ── ROS2 build tools ─────────────────────────────────────────")
        ln("RUN apt-get update && apt-get install -y \\")
        for p in build_apt:
            ln(f"    {p} \\")
        ln("    && rm -rf /var/lib/apt/lists/*")
        if "rosdep" in tools:
            gap()
            ln("RUN {}".format(self._sub(self._cfg["tools"]["rosdep"]["post_install_run"], variables)))
        gap()
        
        if "zsh" in tools:
            ln("# ── Zsh + Oh-My-Zsh ──────────────────────────────────────────")
            ln("RUN apt-get update && apt-get install -y zsh && rm -rf /var/lib/apt/lists/*")
            ln(f"RUN {self._cfg['tools']['zsh']['post_install_run']}")
            gap()
            
        ros_apt = []
        for pkg_key in packages:
            ros_apt.extend(self.resolve_ros_package_apt(pkg_key, distro))
        if ros_apt:
            ln("# ── ROS2 packages ────────────────────────────────────────────")
            ln("RUN apt-get update && apt-get install -y \\")
            for p in ros_apt:
                ln(f"    {p} \\")
            ln("    && rm -rf /var/lib/apt/lists/*")
            gap()
            
        for pkg_key in packages:
            pkg = self._cfg["ros_packages"].get(pkg_key)
            if pkg and pkg.get("env"):
                for k, v in pkg["env"].items():
                    ln(f"ENV {k}={v}")
                gap()
                
        if "ssh" in tools:
            ssh = self._cfg["tools"]["ssh"]
            ln("# ── SSH server ────────────────────────────────────────────────")
            ln("RUN apt-get update && apt-get install -y {} && \\".format(" ".join(ssh["apt"])))
            ln(f"    {ssh['post_install_run']} && \\")
            ln("    rm -rf /var/lib/apt/lists/*")
            ln(f"EXPOSE {ssh['expose_port']}")
            gap()
            
        if not is_root:
            ln("# ── Non-root user ────────────────────────────────────────────")
            ln(f"ARG UID={uid}")
            ln(f"ARG GID={uid}")
            ln("RUN if getent group ${GID} >/dev/null; then \\")
            ln(f"        groupmod -n {username} $(getent group ${{GID}} | cut -d: -f1); \\")
            ln(f"    else groupadd -g ${{GID}} {username}; fi && \\")
            ln("    if getent passwd ${UID} >/dev/null; then \\")
            ln(f"        usermod -l {username} -m -d /home/{username} $(getent passwd ${{UID}} | cut -d: -f1); \\")
            ln(f"    else useradd -m -u ${{UID}} -g ${{GID}} -s /bin/bash {username}; fi")
            if "sudo" in tools:
                ln(f"RUN echo \"{username} ALL=(ALL) NOPASSWD:ALL\" >> /etc/sudoers")
            gap()
            ln(f"RUN mkdir -p {workspace}/src && chown -R {username}:{username} {workspace}")
            gap()
            ln(f"USER {username}")
        else:
            ln(f"RUN mkdir -p {workspace}/src")
            
        ln(f"WORKDIR {workspace}")
        gap()
        
        if "bashrc" in tools:
            home = "/root" if is_root else f"/home/{username}"
            ln("# ── Shell setup ───────────────────────────────────────────────")
            ln(f"RUN echo \"source /opt/ros/{distro}/setup.bash\" >> {home}/.bashrc")
            for pkg_key in packages:
                pkg = self._cfg["ros_packages"].get(pkg_key)
                if pkg and pkg.get("env"):
                    for k, v in pkg["env"].items():
                        ln(f"RUN echo \"export {k}={v}\" >> {home}/.bashrc")
            gap()
            
        for tool_key in tools:
            tool = self._cfg["tools"].get(tool_key)
            if tool and tool.get("env"):
                for k, v in tool["env"].items():
                    ln(f"ENV {k}={v}")
                gap()
                
        if has_cuda:
            ln("# ── NVIDIA GPU runtime ──────────────────────────────────────")
            ln("ENV NVIDIA_VISIBLE_DEVICES=all")
            ln("ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics")
            gap()
            
        ln("CMD [\"/bin/bash\"]")
        return "\n".join(L)

    def build_compose(self, config):
        distro = config["distro"]
        packages = sorted(config["packages"])
        tools = sorted(config["tools"])
        workspace = config["workspace"]
        container_name = config.get("containerName", config.get("container_name"))
        user_type = config.get("userType", config.get("user_type"))
        
        is_root = user_type == "root"
        has_cuda = "cuda" in packages or "tensorrt" in packages
        
        L = []
        def ln(s): L.append(s)
        
        ln(f"# docker-compose.yml — generated by ros2-dockergen v{self._cfg['version']}")
        ln("services:")
        ln(f"  {container_name}:")
        ln("    build:")
        ln("      context: .")
        ln("      dockerfile: Dockerfile")
        if not is_root:
            ln("      args:")
            ln("        UID: ${UID:-1000}")
            ln("        GID: ${GID:-1000}")
        ln(f"    image: ros2-{distro}-{container_name}:latest")
        ln(f"    container_name: {container_name}")
        ln(f"    hostname: {container_name}")
        ln("    stdin_open: true")
        ln("    tty: true")
        ln("    restart: unless-stopped")
        ln("    network_mode: host")
        ln("    environment:")
        ln(f"      - ROS_DISTRO={distro}")
        
        for tool_key in tools:
            tool = self._cfg["tools"].get(tool_key)
            if tool and tool.get("compose_env"):
                for k, v in tool["compose_env"].items():
                    ln(f"      - {k}={v}")
                    
        for pkg_key in packages:
            pkg = self._cfg["ros_packages"].get(pkg_key)
            if pkg and pkg.get("env"):
                for k, v in pkg["env"].items():
                    ln(f"      - {k}={v}")
                    
        ln("    volumes:")
        ln(f"      - ./ros2_ws:{workspace}:rw")
        for tool_key in tools:
            tool = self._cfg["tools"].get(tool_key)
            if tool and tool.get("compose_volumes"):
                for v in tool["compose_volumes"]:
                    ln(f"      - {v}")
                    
        if has_cuda:
            ln("    deploy:")
            ln("      resources:")
            ln("        reservations:")
            ln("          devices:")
            ln("            - driver: nvidia")
            ln("              count: all")
            ln("              capabilities: [gpu, compute, utility]")
            
        ports = []
        for tool_key in tools:
            tool = self._cfg["tools"].get(tool_key)
            if tool and tool.get("compose_port"):
                ports.append(tool["compose_port"])
        if ports:
            ln("    ports:")
            for p in ports:
                ln(f"      - \"{p}\"")
                
        return "\n".join(L)

    def build_readme(self, config):
        distro = config["distro"]
        variant = config["variant"]
        packages = sorted(config["packages"])
        tools = sorted(config["tools"])
        username = config["username"]
        workspace = config["workspace"]
        container_name = config.get("containerName", config.get("container_name"))
        user_type = config.get("userType", config.get("user_type"))
        
        u = "root" if user_type == "root" else username
        d = self._cfg["distros"][distro]
        
        return f"""# ROS2 {d['label']} Docker Environment

Generated by **ros2-dockergen v{self._cfg['version']}**

## What's Inside
| | |
|---|---|
| **ROS2 Distro** | {distro} ({variant}) |
| **Packages** | {", ".join(sorted(packages)) if packages else "none"} |
| **Dev Tools** | {", ".join(sorted(tools))} |
| **User** | {u} |
| **Workspace** | {workspace} |

## Prerequisites
- Docker Engine ≥ 24 / Docker Desktop
- docker compose v2

## Quick Start
```bash
docker compose build
docker compose up -d
docker exec -it {container_name} bash

cd {workspace}
colcon build --symlink-install
source install/setup.bash

docker compose down
```

---
*Generated by ros2-dockergen v{self._cfg['version']}*
"""

    def get_distros(self):
        return [{"value": k, "label": v["label"], "ubuntu": v["ubuntu"], "recommended": v.get("recommended", False)} 
                for k, v in self._cfg["distros"].items()]

    def get_variants(self):
        return [{"value": k, "label": v["label"], "description": v["description"]} 
                for k, v in self._cfg["variants"].items()]

    def get_ros_package_choices(self):
        return [{"value": k, "label": v["label"], "description": v["description"]} 
                for k, v in self._cfg["ros_packages"].items()]

    def get_tool_choices(self):
        return [{"value": k, "label": v["label"], "description": v["description"], "default": v.get("default", False)} 
                for k, v in self._cfg["tools"].items()]
