stackpilot/backend/services/template_service.py

"""Template library — bundled (on-disk) + custom (DB)."""
from __future__ import annotations

import json
import os
import re
from functools import lru_cache
from typing import Optional

from sqlmodel import Session, select

from models.template import Template

_TEMPLATES_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "templates")
_VAR_RE = re.compile(r"\{\{\s*([A-Za-z0-9_]+)\s*\}\}")


@lru_cache(maxsize=1)
def _manifest() -> list[dict]:
    path = os.path.join(_TEMPLATES_DIR, "manifest.json")
    try:
        with open(path, "r", encoding="utf-8") as fh:
            return json.load(fh)
    except (OSError, json.JSONDecodeError):
        return []


def _read_template_file(filename: str) -> str:
    path = os.path.join(_TEMPLATES_DIR, filename)
    try:
        with open(path, "r", encoding="utf-8") as fh:
            return fh.read()
    except OSError:
        return ""


def extract_variables(yaml_str: str) -> list[str]:
    seen: list[str] = []
    for m in _VAR_RE.finditer(yaml_str):
        if m.group(1) not in seen:
            seen.append(m.group(1))
    return seen


def render(yaml_str: str, values: dict[str, str]) -> str:
    def repl(m: re.Match) -> str:
        key = m.group(1)
        return str(values.get(key, m.group(0)))

    return _VAR_RE.sub(repl, yaml_str)


# --------------------------------------------------------------------------- #
# Listing
# --------------------------------------------------------------------------- #


def list_templates(session: Session) -> list[dict]:
    out: list[dict] = []
    for entry in _manifest():
        out.append(
            {
                "id": entry["id"],
                "name": entry["name"],
                "description": entry.get("description"),
                "tags": entry.get("tags", []),
                "gpu": entry.get("gpu"),
                "source": "bundled",
            }
        )
    for tpl in session.exec(select(Template)).all():
        out.append(
            {
                "id": f"custom:{tpl.slug}",
                "name": tpl.name,
                "description": tpl.description,
                "tags": [t for t in tpl.tags.split(",") if t],
                "gpu": None,
                "source": "custom",
            }
        )
    return out


def get_template(session: Session, template_id: str) -> Optional[dict]:
    if template_id.startswith("custom:"):
        slug = template_id.split(":", 1)[1]
        tpl = session.exec(select(Template).where(Template.slug == slug)).first()
        if not tpl:
            return None
        variables = [
            {"name": v, "description": "", "default": ""}
            for v in extract_variables(tpl.yaml)
        ]
        return {
            "id": template_id,
            "name": tpl.name,
            "description": tpl.description,
            "tags": [t for t in tpl.tags.split(",") if t],
            "gpu": None,
            "source": "custom",
            "yaml": tpl.yaml,
            "variables": variables,
        }

    for entry in _manifest():
        if entry["id"] == template_id:
            yaml_str = _read_template_file(entry["file"])
            declared = {v["name"]: v for v in entry.get("variables", [])}
            # Merge declared metadata with any vars actually present.
            variables = []
            for name in extract_variables(yaml_str):
                meta = declared.get(name, {})
                variables.append(
                    {
                        "name": name,
                        "description": meta.get("description", ""),
                        "default": meta.get("default", ""),
                    }
                )
            return {
                "id": entry["id"],
                "name": entry["name"],
                "description": entry.get("description"),
                "tags": entry.get("tags", []),
                "gpu": entry.get("gpu"),
                "source": "bundled",
                "yaml": yaml_str,
                "variables": variables,
            }
    return None


def save_custom(
    session: Session, name: str, yaml_str: str, description: str = "", tags: list[str] | None = None
) -> Template:
    slug = re.sub(r"[^a-z0-9_-]+", "-", name.strip().lower()).strip("-") or "template"
    existing = session.exec(select(Template).where(Template.slug == slug)).first()
    if existing:
        existing.name = name
        existing.description = description
        existing.tags = ",".join(tags or [])
        existing.yaml = yaml_str
        session.add(existing)
        session.commit()
        session.refresh(existing)
        return existing
    tpl = Template(
        slug=slug,
        name=name,
        description=description,
        tags=",".join(tags or []),
        yaml=yaml_str,
    )
    session.add(tpl)
    session.commit()
    session.refresh(tpl)
    return tpl


def delete_custom(session: Session, slug: str) -> bool:
    tpl = session.exec(select(Template).where(Template.slug == slug)).first()
    if not tpl:
        return False
    session.delete(tpl)
    session.commit()
    return True