"""
GUARDiA 온프레미스 디자인 스튜디오
디자인 토큰 관리 · 컴포넌트 생성 · CSS/Tailwind 생성 · Ollama 비전 분석
100% 온프레미스 (외부 API 절대 금지)
"""
import os
import httpx
import json
import base64
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
from fastapi import APIRouter, HTTPException, UploadFile, File
from pydantic import BaseModel

router = APIRouter(prefix="/api/design", tags=["Design Studio"])

OLLAMA_BASE  = "http://localhost:11434"
VISION_MODEL = "llava"
CODE_MODEL   = "codellama"
LLM_MODEL    = "llama3"

# 개방망: 외부 비전 API 허용 (GUARDIA_NETWORK_MODE=open)
_OPEN_NET     = os.environ.get("GUARDIA_NETWORK_MODE") == "open"
_EXT_API_KEY  = os.environ.get("OPENAI_API_KEY", "")    # 개방망 전용

# ── 기본 디자인 토큰 ────────────────────────────────────────────────────────
DEFAULT_TOKENS: Dict[str, Any] = {
    "colors": {
        "primary":    "#003366",
        "secondary":  "#005A8C",
        "accent":     "#00A0C8",
        "success":    "#28a745",
        "warning":    "#ffc107",
        "danger":     "#dc3545",
        "info":       "#17a2b8",
        "dark":       "#1a1a2e",
        "light":      "#f8f9fa",
        "bg":         "#0d1117",
        "surface":    "#161b22",
        "border":     "#30363d",
        "text":       "#e6edf3",
        "muted":      "#8b949e",
    },
    "typography": {
        "font_family": "'Noto Sans KR', 'Inter', sans-serif",
        "size_xs":  "0.75rem",
        "size_sm":  "0.875rem",
        "size_md":  "1rem",
        "size_lg":  "1.125rem",
        "size_xl":  "1.25rem",
        "size_2xl": "1.5rem",
        "weight_normal": "400",
        "weight_medium": "500",
        "weight_bold":   "700",
    },
    "spacing": {
        "xs": "0.25rem", "sm": "0.5rem", "md": "1rem",
        "lg": "1.5rem",  "xl": "2rem",   "2xl": "3rem",
    },
    "radius": {
        "sm": "4px", "md": "8px", "lg": "12px", "xl": "16px", "full": "9999px",
    },
    "shadow": {
        "sm":  "0 1px 3px rgba(0,0,0,.3)",
        "md":  "0 4px 6px rgba(0,0,0,.3)",
        "lg":  "0 10px 15px rgba(0,0,0,.3)",
        "glow": "0 0 20px rgba(0,160,200,.3)",
    },
    "transition": {
        "fast":   "150ms ease",
        "normal": "300ms ease",
        "slow":   "500ms ease",
    },
}

# 런타임 토큰 스토어 (DB 없이 인메모리)
_token_store: Dict[str, Any] = dict(DEFAULT_TOKENS)


# ── Pydantic 모델 ──────────────────────────────────────────────────────────

class TokenUpdate(BaseModel):
    category: str
    key: str
    value: str

class CssGenRequest(BaseModel):
    description: str
    framework: str = "css"   # css | tailwind | scss | styled-components
    tokens: Optional[Dict] = None

class ComponentRequest(BaseModel):
    name: str
    description: str
    framework: str = "react"
    variant: str = "default"
    tokens: Optional[Dict] = None

class ScreenAnalysisRequest(BaseModel):
    image_base64: str        # base64 인코딩 이미지
    analysis_type: str = "improve"  # improve | audit | describe | accessibility

class PaletteRequest(BaseModel):
    base_color: str          # hex 색상 (#003366)
    count: int = 5
    mode: str = "analogous"  # analogous | complementary | triadic | monochromatic


# ── 헬퍼 ──────────────────────────────────────────────────────────────────

async def _ollama_text(model: str, prompt: str) -> str:
    async with httpx.AsyncClient(timeout=120.0) as c:
        r = await c.post(
            f"{OLLAMA_BASE}/api/generate",
            json={"model": model, "prompt": prompt, "stream": False,
                  "options": {"temperature": 0.3}},
        )
        if r.status_code != 200:
            raise HTTPException(503, "Ollama 불가")
        return r.json().get("response", "")

async def _ollama_vision(model: str, prompt: str, image_b64: str) -> str:
    """비전 분석 — 개방망에서는 GPT-4o Vision 사용 가능."""
    if _OPEN_NET and _EXT_API_KEY:
        async with httpx.AsyncClient(timeout=120.0) as c:
            r = await c.post(
                "https://api.openai.com/v1/chat/completions",
                headers={"Authorization": f"Bearer {_EXT_API_KEY}"},
                json={
                    "model": "gpt-4o",
                    "messages": [{"role": "user", "content": [
                        {"type": "text", "text": prompt},
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}},
                    ]}],
                    "max_tokens": 1024,
                },
            )
            if r.status_code == 200:
                return r.json()["choices"][0]["message"]["content"]
    # 폐쇄망 fallback — Ollama llava
    async with httpx.AsyncClient(timeout=120.0) as c:
        r = await c.post(
            f"{OLLAMA_BASE}/api/generate",
            json={"model": model, "prompt": prompt,
                  "images": [image_b64], "stream": False},
        )
        if r.status_code != 200:
            raise HTTPException(503, "Ollama vision 불가")
        return r.json().get("response", "")


# ── 토큰 관리 ─────────────────────────────────────────────────────────────

@router.get("/tokens")
async def get_tokens():
    """현재 디자인 토큰 조회."""
    return {"tokens": _token_store, "version": "1.0.0"}

@router.patch("/tokens")
async def update_token(req: TokenUpdate):
    """토큰 값 개별 업데이트."""
    if req.category not in _token_store:
        _token_store[req.category] = {}
    _token_store[req.category][req.key] = req.value
    return {"updated": True, "category": req.category, "key": req.key, "value": req.value}

@router.post("/tokens/reset")
async def reset_tokens():
    """기본 토큰으로 초기화."""
    global _token_store
    _token_store = dict(DEFAULT_TOKENS)
    return {"reset": True}

@router.get("/tokens/export/css")
async def export_css_variables():
    """토큰 → CSS 변수 내보내기."""
    lines = [":root {"]
    for cat, vals in _token_store.items():
        if isinstance(vals, dict):
            for k, v in vals.items():
                lines.append(f"  --{cat}-{k.replace('_', '-')}: {v};")
    lines.append("}")
    return {"css": "\n".join(lines)}

@router.get("/tokens/export/tailwind")
async def export_tailwind_config():
    """토큰 → Tailwind 설정 내보내기."""
    config = {
        "theme": {
            "extend": {
                "colors": _token_store.get("colors", {}),
                "fontFamily": {"sans": [_token_store["typography"].get("font_family", "sans-serif")]},
                "borderRadius": _token_store.get("radius", {}),
                "boxShadow": _token_store.get("shadow", {}),
                "spacing": _token_store.get("spacing", {}),
            }
        }
    }
    return {"config": config, "js": f"module.exports = {json.dumps(config, indent=2, ensure_ascii=False)}"}


# ── CSS 생성 ───────────────────────────────────────────────────────────────

@router.post("/css/generate")
async def generate_css(req: CssGenRequest):
    """자연어 → CSS/Tailwind 생성 (Ollama 온프레미스)."""
    tokens = req.tokens or _token_store
    token_hint = f"디자인 토큰: {json.dumps(tokens.get('colors', {}), ensure_ascii=False)}"

    prompt = (
        f"{req.framework} 스타일로 다음을 구현하는 CSS를 작성하라:\n"
        f"요청: {req.description}\n{token_hint}\n\n"
        "CSS 코드만 출력하라."
    )
    css = await _ollama_text(CODE_MODEL, prompt)
    return {"framework": req.framework, "css": css, "model": CODE_MODEL}


# ── 컴포넌트 생성 ─────────────────────────────────────────────────────────

@router.post("/component/generate")
async def generate_component(req: ComponentRequest):
    """컴포넌트 코드 + CSS 동시 생성."""
    tokens = req.tokens or _token_store

    prompt = (
        f"{req.framework} 컴포넌트를 생성하라.\n"
        f"이름: {req.name}, 설명: {req.description}, 변형: {req.variant}\n"
        f"디자인 토큰: {json.dumps(tokens.get('colors', {}), ensure_ascii=False)}\n\n"
        "완전한 컴포넌트 코드를 출력하라."
    )
    code = await _ollama_text(CODE_MODEL, prompt)

    css_prompt = (
        f"위 컴포넌트({req.name})의 CSS 스타일을 토큰 기반으로 작성하라."
    )
    css = await _ollama_text(CODE_MODEL, css_prompt)

    return {
        "name": req.name,
        "framework": req.framework,
        "variant": req.variant,
        "code": code,
        "css": css,
    }

@router.get("/component/library")
async def list_component_library():
    """GUARDiA 컴포넌트 라이브러리 목록."""
    return {
        "components": [
            {"id": "btn",       "name": "Button",       "variants": ["primary", "secondary", "danger", "ghost"]},
            {"id": "card",      "name": "Card",          "variants": ["default", "elevated", "outlined"]},
            {"id": "table",     "name": "DataTable",     "variants": ["default", "compact", "striped"]},
            {"id": "modal",     "name": "Modal",         "variants": ["default", "large", "fullscreen"]},
            {"id": "form",      "name": "Form",          "variants": ["vertical", "horizontal", "inline"]},
            {"id": "badge",     "name": "Badge",         "variants": ["status", "count", "label"]},
            {"id": "chart",     "name": "Chart",         "variants": ["line", "bar", "donut", "heatmap"]},
            {"id": "timeline",  "name": "Timeline",      "variants": ["vertical", "horizontal"]},
            {"id": "kanban",    "name": "Kanban",        "variants": ["default", "compact"]},
            {"id": "sidebar",   "name": "Sidebar",       "variants": ["dark", "light", "collapsed"]},
        ]
    }


# ── 화면 분석 (Ollama Vision) ─────────────────────────────────────────────

@router.post("/analyze/screen")
async def analyze_screen(req: ScreenAnalysisRequest):
    """스크린샷 → UI 분석 (Ollama llava 비전 모델)."""
    prompts = {
        "improve":       "이 UI 화면을 분석하고 개선 방안을 한국어로 제안하라.",
        "audit":         "이 화면의 UX/접근성 이슈를 체계적으로 감사하라.",
        "describe":      "이 화면의 레이아웃과 컴포넌트를 설명하라.",
        "accessibility": "이 화면의 접근성 문제를 WCAG 2.1 기준으로 평가하라.",
    }
    result = await _ollama_vision(
        VISION_MODEL,
        prompts.get(req.analysis_type, prompts["describe"]),
        req.image_base64,
    )
    return {
        "analysis_type": req.analysis_type,
        "result": result,
        "model": VISION_MODEL,
    }

@router.post("/analyze/upload")
async def analyze_upload(
    file: UploadFile = File(...),
    analysis_type: str = "improve",
):
    """파일 업로드 방식 화면 분석."""
    content = await file.read()
    b64 = base64.b64encode(content).decode()
    prompts = {
        "improve":       "이 UI 화면의 개선 사항을 제안하라.",
        "audit":         "이 화면의 UX 이슈를 감사하라.",
        "describe":      "이 화면의 구성 요소를 설명하라.",
        "accessibility": "이 화면의 접근성을 평가하라.",
    }
    result = await _ollama_vision(VISION_MODEL, prompts.get(analysis_type, "분석하라."), b64)
    return {"filename": file.filename, "analysis_type": analysis_type, "result": result}


# ── 색상 팔레트 ────────────────────────────────────────────────────────────

@router.post("/palette/generate")
async def generate_palette(req: PaletteRequest):
    """기준 색상으로 팔레트 자동 생성."""
    prompt = (
        f"색상 {req.base_color}를 기준으로 {req.mode} 방식의 {req.count}가지 색상 팔레트를 생성하라.\n"
        "JSON 배열로만 출력: [\"#RRGGBB\", ...]"
    )
    result = await _ollama_text(LLM_MODEL, prompt)
    try:
        colors = json.loads(result)
    except json.JSONDecodeError:
        colors = [req.base_color]
    return {"base": req.base_color, "mode": req.mode, "palette": colors}


# ── 디자인 감사 ────────────────────────────────────────────────────────────

@router.post("/audit/tokens")
async def audit_token_usage(component_code: str):
    """컴포넌트 코드가 디자인 토큰을 올바르게 사용하는지 감사."""
    token_keys = [f"--{cat}-{k}".replace("_", "-")
                  for cat, vals in _token_store.items()
                  if isinstance(vals, dict)
                  for k in vals]
    prompt = (
        f"다음 코드에서 하드코딩된 색상/크기/간격을 찾고 적절한 CSS 변수로 교체 제안하라.\n"
        f"사용 가능한 토큰: {token_keys[:20]}\n\n{component_code}\n\n"
        "JSON 출력: {\"issues\": [{\"line\": N, \"original\": \"...\", \"suggested\": \"...\"}], \"score\": N}"
    )
    result = await _ollama_text(LLM_MODEL, prompt)
    try:
        return {"audit": json.loads(result)}
    except json.JSONDecodeError:
        return {"audit": {"issues": [], "score": 100, "raw": result}}

@router.get("/health")
async def design_health():
    """디자인 스튜디오 상태."""
    try:
        async with httpx.AsyncClient(timeout=5.0) as c:
            r = await c.get(f"{OLLAMA_BASE}/api/tags")
            models = [m["name"] for m in r.json().get("models", [])]
        return {
            "status": "healthy",
            "vision_model": {"name": VISION_MODEL, "available": any(VISION_MODEL in m for m in models)},
            "code_model":   {"name": CODE_MODEL,   "available": any(CODE_MODEL in m for m in models)},
            "token_categories": list(_token_store.keys()),
        }
    except Exception as e:
        return {"status": "down", "error": str(e)}
