guardia-itsm/routers/icon_generator.py

"""SVG 아이콘 생성 — 텍스트 설명→SVG 코드 자동 생성"""
from __future__ import annotations
import json, logging, re
from datetime import datetime
from pathlib import Path
from typing import Optional

import httpx
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel
from sqlalchemy import select, desc, func
from sqlalchemy.ext.asyncio import AsyncSession

from core.auth import get_current_user
from database import get_db
from models import User, GeneratedIcon

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/icon", tags=["아이콘 생성"])
OLLAMA_URL = "http://localhost:11434"
TEXT_MODEL = "llama3"

# 내장 아이콘 템플릿 (GUARDiA 브랜드 #003366)
BUILTIN_ICONS = {
    "server": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="#003366"><rect x="2" y="3" width="20" height="6" rx="1"/><rect x="2" y="11" width="20" height="6" rx="1"/><circle cx="19" cy="6" r="1.5" fill="white"/><circle cx="19" cy="14" r="1.5" fill="white"/></svg>',
    "alert":  '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="#003366"><path d="M12 2L2 20h20L12 2z"/><text x="11" y="17" font-size="9" fill="white">!</text></svg>',
    "check":  '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="#003366" stroke-width="2.5"><polyline points="20 6 9 17 4 12"/></svg>',
    "deploy": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="#003366"><path d="M12 2l8 4v6c0 5-4 9-8 10C8 21 4 17 4 12V6l8-4z"/></svg>',
    "user":   '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="#003366"><circle cx="12" cy="8" r="4"/><path d="M4 20c0-4 3.6-7 8-7s8 3 8 7"/></svg>',
    "database": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="#003366"><ellipse cx="12" cy="5" rx="8" ry="3"/><path d="M4 5v14c0 1.7 3.6 3 8 3s8-1.3 8-3V5"/><path d="M4 12c0 1.7 3.6 3 8 3s8-1.3 8-3" fill="none" stroke="white" stroke-width="1"/></svg>',
    "network": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="#003366"><circle cx="12" cy="12" r="3"/><circle cx="4" cy="6" r="2"/><circle cx="20" cy="6" r="2"/><circle cx="4" cy="18" r="2"/><circle cx="20" cy="18" r="2"/><line x1="12" y1="9" x2="4" y2="6" stroke="#003366" stroke-width="1.5"/><line x1="12" y1="9" x2="20" y2="6" stroke="#003366" stroke-width="1.5"/><line x1="12" y1="15" x2="4" y2="18" stroke="#003366" stroke-width="1.5"/><line x1="12" y1="15" x2="20" y2="18" stroke="#003366" stroke-width="1.5"/></svg>',
    "security": '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="#003366"><path d="M12 2l8 4v5c0 5.5-3.5 10-8 11C7.5 21 4 16.5 4 11V6l8-4z"/><path d="M9 12l2 2 4-4" fill="none" stroke="white" stroke-width="2"/></svg>',
}

SVG_SYSTEM_PROMPT = """당신은 SVG 아이콘 전문 디자이너입니다.
요구사항에 맞는 24x24 픽토그램 SVG 코드를 생성합니다.

규칙:
1. viewBox="0 0 24 24" 사용
2. xmlns="http://www.w3.org/2000/svg" 포함
3. 기본 색상: fill="{color}"
4. 단순하고 명확한 픽토그램 스타일
5. SVG 코드만 출력 (다른 텍스트 없음)
6. 반드시 완전한 SVG 태그로 시작하고 끝내기"""


async def _generate_svg(description: str, color: str = "#003366", size: int = 24) -> str:
    """Ollama로 SVG 아이콘 생성."""
    system = SVG_SYSTEM_PROMPT.replace("{color}", color)
    prompt = f"다음 아이콘을 SVG로 생성 (viewBox 0 0 {size} {size}): {description}"

    try:
        async with httpx.AsyncClient(timeout=30) as c:
            r = await c.post(f"{OLLAMA_URL}/api/generate", json={
                "model": TEXT_MODEL,
                "system": system,
                "prompt": prompt,
                "stream": False,
            })
            resp = r.json().get("response", "")

        # SVG 추출
        svg_match = re.search(r'(<svg[\s\S]+?</svg>)', resp, re.IGNORECASE)
        if svg_match:
            return svg_match.group(1)

        # 코드 블록에서 추출
        block_match = re.search(r'```(?:svg|xml)?\n?([\s\S]+?)```', resp)
        if block_match:
            candidate = block_match.group(1).strip()
            if '<svg' in candidate:
                return candidate

        # 폴백: 기본 아이콘
        return _default_icon(description, color, size)
    except Exception as e:
        logger.warning(f"SVG 생성 실패: {e}")
        return _default_icon(description, color, size)


def _default_icon(name: str, color: str, size: int) -> str:
    """폴백: 텍스트 이니셜 SVG."""
    initial = name[0].upper() if name else "?"
    return (f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {size} {size}">'
            f'<circle cx="{size//2}" cy="{size//2}" r="{size//2-1}" fill="{color}"/>'
            f'<text x="{size//2}" y="{size//2+4}" text-anchor="middle" '
            f'font-size="{size//2}" fill="white" font-family="sans-serif">{initial}</text></svg>')


class IconRequest(BaseModel):
    description: str
    color: str = "#003366"
    size: int = 24
    category: str = "custom"
    save: bool = True


class BatchRequest(BaseModel):
    icons: list[IconRequest]


@router.post("/generate", status_code=201)
async def generate_icon(body: IconRequest, db: AsyncSession = Depends(get_db),
                        user: User = Depends(get_current_user)):
    """텍스트 설명→SVG 아이콘 생성."""
    # 내장 아이콘 확인
    for builtin_name, builtin_svg in BUILTIN_ICONS.items():
        if builtin_name in body.description.lower():
            return {"icon_id": f"builtin-{builtin_name}", "name": builtin_name,
                    "svg_code": builtin_svg, "source": "builtin",
                    "download_url": f"/api/icon/builtin/{builtin_name}"}

    svg = await _generate_svg(body.description, body.color, body.size)

    if not body.save:
        return {"svg_code": svg, "source": "generated"}

    icon = GeneratedIcon(
        name=body.description[:50], description=body.description,
        svg_code=svg, category=body.category,
        primary_color=body.color, size=body.size,
        created_by=user.id, created_at=datetime.utcnow(),
    )
    db.add(icon); await db.commit(); await db.refresh(icon)

    return {
        "icon_id": icon.id, "name": icon.name, "svg_code": svg,
        "source": "ai_generated", "model": TEXT_MODEL,
        "download_url": f"/api/icon/{icon.id}/download",
    }


@router.post("/batch")
async def batch_generate(body: BatchRequest, db: AsyncSession = Depends(get_db),
                         user: User = Depends(get_current_user)):
    """여러 아이콘 일괄 생성."""
    results = []
    for req in body.icons[:10]:  # 최대 10개
        svg = await _generate_svg(req.description, req.color, req.size)
        icon = GeneratedIcon(
            name=req.description[:50], description=req.description,
            svg_code=svg, category=req.category, primary_color=req.color, size=req.size,
            created_by=user.id, created_at=datetime.utcnow(),
        )
        db.add(icon)
        results.append({"description": req.description, "svg_preview": svg[:100] + "..."})
    await db.commit()
    return {"generated": len(results), "icons": results}


@router.get("/library")
async def icon_library(db: AsyncSession = Depends(get_db), user: User = Depends(get_current_user)):
    # 내장 아이콘
    builtins = [{"id": f"builtin-{k}", "name": k, "source": "builtin",
                 "download_url": f"/api/icon/builtin/{k}"} for k in BUILTIN_ICONS]
    # 생성 아이콘
    rows = await db.execute(
        select(GeneratedIcon).order_by(desc(GeneratedIcon.created_at)).limit(50)
    )
    custom = [{"id": ic.id, "name": ic.name, "category": ic.category,
               "source": "ai_generated", "download_url": f"/api/icon/{ic.id}/download"}
              for ic in rows.scalars().all()]
    return {"builtin": builtins, "custom": custom, "total": len(builtins) + len(custom)}


@router.get("/builtin/{name}")
async def get_builtin(name: str):
    """내장 아이콘 SVG 반환."""
    if name not in BUILTIN_ICONS:
        raise HTTPException(404, f"내장 아이콘 '{name}' 없음")
    return Response(content=BUILTIN_ICONS[name], media_type="image/svg+xml",
                    headers={"Content-Disposition": f'attachment; filename="{name}.svg"'})


@router.get("/{icon_id}/download")
async def download_icon(icon_id: int, db: AsyncSession = Depends(get_db),
                        user: User = Depends(get_current_user)):
    row = await db.execute(select(GeneratedIcon).where(GeneratedIcon.id == icon_id))
    icon = row.scalar_one_or_none()
    if not icon: raise HTTPException(404)
    return Response(content=icon.svg_code, media_type="image/svg+xml",
                    headers={"Content-Disposition": f'attachment; filename="{icon.name}.svg"'})


@router.post("/customize/{icon_id}")
async def customize_icon(icon_id: int, color: str = "#003366", size: int = 24,
                         db: AsyncSession = Depends(get_db), user: User = Depends(get_current_user)):
    """색상·크기 변경."""
    row = await db.execute(select(GeneratedIcon).where(GeneratedIcon.id == icon_id))
    icon = row.scalar_one_or_none()
    if not icon: raise HTTPException(404)
    customized = re.sub(r'fill="[^"]*"', f'fill="{color}"', icon.svg_code)
    return {"svg_code": customized, "color": color, "size": size}