server/main.py

"""Campaign Concept Studio — FastAPI backend.

Exposes three endpoints:
  * POST /api/generate        — generate a structured campaign concept via OpenAI Responses API
  * POST /api/generate-image  — generate a 1024x1024 image via OpenAI Images API
  * GET  /api/health          — health check

All OpenAI calls happen server-side; the frontend never sees the API key.
Uses httpx.AsyncClient (async) to avoid blocking the event loop.
"""

from __future__ import annotations

import os
import json
import asyncio
import logging
from pathlib import Path

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")

RESPONSES_URL = "https://api.openai.com/v1/responses"
IMAGES_URL = "https://api.openai.com/v1/images/generations"

# Upgraded models (Sept 2026):
#   gpt-5.6-luna = cost-sensitive flagship, replaces gpt-4.1
#   gpt-image-2 = current image model, replaces deprecated gpt-image-1
TEXT_MODEL = os.environ.get("OPENAI_TEXT_MODEL", "gpt-5.6-luna")
IMAGE_MODEL = os.environ.get("OPENAI_IMAGE_MODEL", "gpt-image-2")

BASE_DIR = Path(__file__).resolve().parent.parent
FRONTEND_DIR = BASE_DIR / "frontend"
INDEX_FILE = FRONTEND_DIR / "index.html"

logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(levelname)-7s  %(message)s")
logger = logging.getLogger("campaign-studio")

# ---------------------------------------------------------------------------
# Request / response models
# ---------------------------------------------------------------------------


class GenerateRequest(BaseModel):
    brief: str = Field(..., min_length=1, description="Campaign brief / goal")
    audience: str = Field(default="", description="Target audience")
    product: str = Field(default="", description="Product / service details")
    tone: str = Field(default="professional", description="Desired tone of voice")
    channels: list[str] = Field(default_factory=list, description="Selected channels")


class GenerateImageRequest(BaseModel):
    prompt: str = Field(..., min_length=1, description="Image prompt")


# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(title="Campaign Concept Studio", version="2.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Serve static frontend assets (if any beyond index.html) from /frontend
if FRONTEND_DIR.is_dir():
    app.mount("/static", StaticFiles(directory=str(FRONTEND_DIR)), name="static")


# ---------------------------------------------------------------------------
# Error mapping (user-friendly, no API detail leaks)
# ---------------------------------------------------------------------------

_ERROR_MAP = {
    401: "Service not configured. Please contact support.",
    403: "Access denied. Please contact support.",
    429: "Service is busy. Please try again in a moment.",
    500: "AI service temporarily unavailable. Please try again.",
    502: "AI service temporarily unavailable. Please try again.",
    503: "AI service temporarily unavailable. Please try again.",
}


def _friendly_error(status_code: int, detail: str = "") -> HTTPException:
    """Return an HTTPException with a user-friendly message, logging the real error."""
    safe_msg = _ERROR_MAP.get(status_code, f"Request failed ({status_code}). Please try again.")
    if detail:
        logger.error("OpenAI API error %d: %s", status_code, detail[:300])
    return HTTPException(status_code=502 if status_code >= 500 else status_code, detail=safe_msg)


def _safe_error_detail(resp: httpx.Response) -> str:
    """Extract human-readable error from an OpenAI error response (for logging only)."""
    try:
        body = resp.json()
        err = body.get("error") or body
        if isinstance(err, dict):
            return err.get("message", json.dumps(err))
        return str(err)
    except Exception:
        return resp.text[:500]


# ---------------------------------------------------------------------------
# OpenAI helpers (async with retry)
# ---------------------------------------------------------------------------

MAX_RETRIES = 3
RETRY_DELAYS = [1, 2, 4]  # seconds
RETRYABLE_STATUS = {429, 500, 502, 503}

# JSON schema for structured output (strict mode guarantees all fields present)
CAMPAIGN_SCHEMA = {
    "type": "object",
    "properties": {
        "concept": {
            "type": "string",
            "description": "A 2-4 sentence campaign concept.",
        },
        "copy_variants": {
            "type": "array",
            "minItems": 3,
            "maxItems": 3,
            "items": {
                "type": "object",
                "properties": {
                    "headline": {"type": "string"},
                    "body": {"type": "string"},
                },
                "required": ["headline", "body"],
                "additionalProperties": False,
            },
        },
        "launch_checklist": {
            "type": "array",
            "minItems": 5,
            "items": {"type": "string"},
        },
        "image_prompts": {
            "type": "array",
            "minItems": 3,
            "items": {"type": "string"},
        },
    },
    "required": ["concept", "copy_variants", "launch_checklist", "image_prompts"],
    "additionalProperties": False,
}


def _openai_headers() -> dict[str, str]:
    if not OPENAI_API_KEY:
        raise HTTPException(
            status_code=500,
            detail="Service not configured. Please contact support.",
        )
    return {
        "Authorization": f"Bearer {OPENAI_API_KEY}",
        "Content-Type": "application/json",
    }


def _build_instruction(audience: str, product: str, tone: str, channels: list[str]) -> str:
    """Build the system-style instruction for the Responses API."""
    channel_txt = ", ".join(channels) if channels else "unspecified"
    parts = [
        "You are a senior creative director at a top marketing agency.",
        "Generate a complete campaign concept based on the user brief.",
        "Return ONLY valid JSON with this exact shape:",
        "{",
        '  "concept": "<2-4 sentence campaign concept>",',
        '  "copy_variants": [',
        '    {"headline": "<short punchy headline>", "body": "<1-3 sentence supporting body copy>"},',
        '    {"headline": "<short punchy headline>", "body": "<1-3 sentence supporting body copy>"},',
        '    {"headline": "<short punchy headline>", "body": "<1-3 sentence supporting body copy>"}',
        "  ],",
        '  "launch_checklist": ["<actionable checklist item>", "..."],',
        '  "image_prompts": ["<detailed image generation prompt>", "<...>", "<...>"]',
        "}",
        "Produce exactly 3 copy_variants, at least 5 launch_checklist items, and 3 image_prompts.",
        "Make the concept distinctive, memorable, and aligned with the brief.",
    ]
    if audience:
        parts.append(f"Target audience: {audience}.")
    if product:
        parts.append(f"Product / service: {product}.")
    if tone:
        parts.append(f"Desired tone of voice: {tone}.")
    if channels:
        parts.append(f"Channels to plan for: {channel_txt}.")
    return "\n".join(parts)


async def _retry_request(
    method: str,
    url: str,
    headers: dict,
    json_payload: dict | None = None,
    content: str | None = None,
    timeout: float = 90.0,
) -> httpx.Response:
    """Execute an HTTP request with exponential backoff retry on transient failures."""
    last_exc = None
    last_resp = None

    for attempt in range(MAX_RETRIES):
        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                if content is not None:
                    resp = await client.post(url, headers=headers, content=content)
                else:
                    resp = await client.post(url, headers=headers, json=json_payload)
            if resp.status_code not in RETRYABLE_STATUS:
                return resp
            last_resp = resp
            logger.warning(
                "OpenAI returned %d (attempt %d/%d), retrying in %ds",
                resp.status_code, attempt + 1, MAX_RETRIES, RETRY_DELAYS[attempt],
            )
        except httpx.HTTPError as exc:
            last_exc = exc
            logger.warning("Network error (attempt %d/%d): %s", attempt + 1, MAX_RETRIES, exc)

        if attempt < MAX_RETRIES - 1:
            await asyncio.sleep(RETRY_DELAYS[attempt])

    # All retries exhausted
    if last_resp:
        raise _friendly_error(last_resp.status_code, _safe_error_detail(last_resp))
    raise HTTPException(status_code=502, detail="Could not reach AI service. Please try again.")


async def call_responses_api(user_brief: str, instruction: str) -> dict:
    """Call the OpenAI Responses API and return the parsed JSON payload."""
    payload = {
        "model": TEXT_MODEL,
        "instructions": instruction,
        "input": user_brief,
        "max_output_tokens": 1200,
        # Use json_schema strict mode — guarantees all fields present
        "text": {
            "format": {
                "type": "json_schema",
                "name": "campaign_concept",
                "strict": True,
                "schema": CAMPAIGN_SCHEMA,
            }
        },
        # Prompt caching — saves ~50% on repeated system prompts
        "prompt_cache_key": "campaign-studio-v2",
    }

    try:
        resp = await _retry_request(
            "POST", RESPONSES_URL, _openai_headers(), json_payload=payload, timeout=90.0
        )
    except HTTPException:
        raise

    if resp.status_code != 200:
        raise _friendly_error(resp.status_code, _safe_error_detail(resp))

    data = resp.json()

    # The Responses API returns output[] -> content[] -> text for text output.
    text = ""
    try:
        for item in data.get("output", []):
            for block in item.get("content", []):
                if block.get("type") in ("output_text", "text"):
                    text += block.get("text", "")
    except (AttributeError, TypeError):
        pass

    if not text:
        text = data.get("output_text", "")

    if not text:
        raise HTTPException(
            status_code=502,
            detail="AI service returned an empty response. Please try again.",
        )

    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        # Sometimes the model wraps JSON in stray quotes/backticks
        cleaned = text.strip().strip("`")
        if cleaned.startswith("json"):
            cleaned = cleaned[4:].strip()
        try:
            parsed = json.loads(cleaned)
        except json.JSONDecodeError:
            raise HTTPException(
                status_code=502,
                detail="AI service returned an unexpected format. Please try again.",
            )

    # Validate the expected structure (json_schema strict mode should guarantee this)
    if not isinstance(parsed, dict) or "concept" not in parsed:
        raise HTTPException(
            status_code=502,
            detail="AI service returned an unexpected format. Please try again.",
        )

    return parsed


async def call_images_api(prompt: str) -> str:
    """Call the OpenAI Images API and return a base64-encoded PNG string."""
    payload = {
        "model": IMAGE_MODEL,
        "prompt": prompt,
        "n": 1,
        "size": "1024x1024",
    }

    try:
        resp = await _retry_request(
            "POST", IMAGES_URL, _openai_headers(), json_payload=payload, timeout=120.0
        )
    except HTTPException:
        raise

    if resp.status_code != 200:
        raise _friendly_error(resp.status_code, _safe_error_detail(resp))

    data = resp.json()
    try:
        b64 = data["data"][0]["b64_json"]
    except (KeyError, IndexError, TypeError):
        # Some image models return a url instead; fall back gracefully
        url = None
        try:
            url = data["data"][0]["url"]
        except (KeyError, IndexError, TypeError):
            pass
        if url:
            try:
                async with httpx.AsyncClient(timeout=60.0) as client:
                    img_resp = await client.get(url)
                    if img_resp.status_code == 200:
                        import base64
                        b64 = base64.b64encode(img_resp.content).decode("ascii")
                        return b64
            except httpx.HTTPError:
                pass
        raise HTTPException(
            status_code=502,
            detail="AI image service returned an unexpected response. Please try again.",
        )

    return b64


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@app.get("/api/health")
async def health() -> dict:
    return {"status": "ok"}


@app.post("/api/generate")
async def generate(req: GenerateRequest) -> dict:
    instruction = _build_instruction(req.audience, req.product, req.tone, req.channels)
    result = await call_responses_api(req.brief, instruction)
    return result


@app.post("/api/generate-image")
async def generate_image(req: GenerateImageRequest) -> dict:
    b64 = await call_images_api(req.prompt)
    return {"url": f"data:image/png;base64,{b64}"}


@app.get("/")
async def root() -> FileResponse:
    if not INDEX_FILE.is_file():
        raise HTTPException(status_code=404, detail="frontend/index.html not found")
    return FileResponse(str(INDEX_FILE))


# ---------------------------------------------------------------------------
# Catch-all for unknown /api routes (returns JSON, not HTML)
# ---------------------------------------------------------------------------


@app.exception_handler(404)
async def not_found_handler(request, exc):  # type: ignore[no-untyped-def]
    if request.url.path.startswith("/api"):
        return JSONResponse(status_code=404, content={"detail": "Not found"})
    if INDEX_FILE.is_file() and not request.url.path.startswith("/static"):
        return FileResponse(str(INDEX_FILE))
    return JSONResponse(status_code=404, content={"detail": "Not found"})


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("server.main:app", host="0.0.0.0", port=8000, reload=True)