cadSet/CADDesigner-Code/web_interface/artifacts.py

from __future__ import annotations

from pathlib import Path
import mimetypes
import re
from typing import Any, Dict, Iterable, Optional


PROJECT_ROOT = Path(__file__).resolve().parents[1]
CODE_FILE_PATTERN = re.compile(r"<\|code_file\|>([^<]*?)</\|code_file(?:\|)?>")
OUTPUT_FILE_PATTERN = re.compile(r"<\|output_file\|>([^<]*?)</\|output_file(?:\|)?>")
WORKSPACE_DIRNAME = "workspace"

MODEL_CONTENT_TYPES = {
    ".stl": "model/stl",
    ".obj": "text/plain; charset=utf-8",
    ".ply": "application/octet-stream",
    ".glb": "model/gltf-binary",
    ".gltf": "model/gltf+json",
    ".step": "model/step",
    ".stp": "model/step",
}


def content_to_text(content: Any) -> str:
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        fragments: list[str] = []
        for item in content:
            if isinstance(item, dict):
                item_type = item.get("type")
                if item_type == "text":
                    fragments.append(str(item.get("text", "")))
                elif item_type == "image_url":
                    fragments.append("[Image]")
                else:
                    fragments.append(str(item))
            else:
                fragments.append(str(item))
        return "\n".join(fragment for fragment in fragments if fragment).strip()
    return str(content)


def resolve_project_path(
    path_text: str, project_root: Path = PROJECT_ROOT
) -> Optional[Path]:
    candidate_text = path_text.strip()
    if not candidate_text:
        return None

    candidate = Path(candidate_text).expanduser()
    if not candidate.is_absolute():
        candidate = project_root / candidate

    try:
        resolved_project_root = project_root.resolve(strict=False)
        resolved = candidate.resolve(strict=False)
        resolved.relative_to(resolved_project_root)
    except (OSError, ValueError):
        return None

    return resolved


def resolve_artifact_path(
    path_text: str,
    project_root: Path = PROJECT_ROOT,
    reference_path: Optional[Path] = None,
) -> Optional[Path]:
    candidate_text = path_text.strip()
    if not candidate_text:
        return None

    candidate_paths: list[Path] = []
    direct_path = resolve_project_path(candidate_text, project_root=project_root)
    if direct_path is not None:
        candidate_paths.append(direct_path)

    raw_candidate = Path(candidate_text).expanduser()
    if raw_candidate.is_absolute():
        return (
            direct_path if direct_path is not None and direct_path.is_file() else None
        )

    normalized_relative = candidate_text.replace("\\", "/")
    while normalized_relative.startswith("./"):
        normalized_relative = normalized_relative[2:]

    relative_fragments = [candidate_text]
    if normalized_relative and normalized_relative != candidate_text:
        relative_fragments.append(normalized_relative)

    search_roots: list[Path] = []
    if reference_path is not None:
        search_roots.append(reference_path.parent)

    workspace_root = project_root / WORKSPACE_DIRNAME
    if workspace_root.exists():
        search_roots.append(workspace_root)

    try:
        seen: set[Path] = set(candidate_paths)
        for root in search_roots:
            for fragment in relative_fragments:
                resolved = resolve_project_path(
                    str(root / fragment), project_root=project_root
                )
                if resolved is not None and resolved not in seen:
                    candidate_paths.append(resolved)
                    seen.add(resolved)
    except Exception:
        return None

    for candidate_path in candidate_paths:
        if candidate_path.is_file():
            return candidate_path

    return None


def remember_recent_path(paths: list[Path], next_path: Path) -> None:
    try:
        paths.remove(next_path)
    except ValueError:
        pass

    paths.append(next_path)


def paths_by_recency(paths: list[Path]) -> list[Path]:
    return list(reversed(paths))


def guess_content_type(path: Path) -> str:
    suffix = path.suffix.lower()
    explicit = MODEL_CONTENT_TYPES.get(suffix)
    if explicit:
        return explicit

    guessed, _ = mimetypes.guess_type(path.name)
    return guessed or "application/octet-stream"


def latest_stl_for_code_file(code_path: Path) -> Optional[Path]:
    if code_path.name != "model.py" or not code_path.exists():
        return None

    stl_files = [path for path in code_path.parent.glob("*.stl") if path.is_file()]
    if not stl_files:
        return None

    return max(stl_files, key=lambda item: item.stat().st_mtime)


def extract_latest_artifacts(
    messages: Iterable[Dict[str, Any]], project_root: Path = PROJECT_ROOT
) -> Dict[str, Any]:
    latest_code_path: Optional[Path] = None
    latest_model_path: Optional[Path] = None
    code_paths: list[Path] = []
    model_paths: list[Path] = []
    output_file_paths: list[Path] = []

    for message in messages:
        if message.get("role") != "assistant":
            continue

        text = content_to_text(message.get("content"))

        for match in CODE_FILE_PATTERN.findall(text):
            resolved = resolve_artifact_path(match, project_root=project_root)
            if resolved is not None:
                latest_code_path = resolved
                remember_recent_path(code_paths, resolved)

        for match in OUTPUT_FILE_PATTERN.findall(text):
            resolved = resolve_artifact_path(
                match,
                project_root=project_root,
                reference_path=latest_code_path,
            )
            if resolved is not None:
                remember_recent_path(output_file_paths, resolved)
                latest_model_path = resolved
                remember_recent_path(model_paths, resolved)

    if latest_code_path is not None and latest_model_path is None:
        latest_model_path = latest_stl_for_code_file(latest_code_path)
        if latest_model_path is not None:
            remember_recent_path(model_paths, latest_model_path)

    return {
        "code_path": latest_code_path,
        "code_paths": paths_by_recency(code_paths),
        "model_path": latest_model_path,
        "model_paths": paths_by_recency(model_paths),
        "output_paths": paths_by_recency(output_file_paths),
    }


__all__ = [
    "PROJECT_ROOT",
    "content_to_text",
    "extract_latest_artifacts",
    "guess_content_type",
    "resolve_artifact_path",
    "resolve_project_path",
]