api.py

import os
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel
from typing import List, Optional, Any

from modules.config import FEATURE_SEQUENCE, SECTIONS
from modules.core_logic import (
    generate_prompt as core_generate_prompt,
    handle_regeneration as core_handle_regeneration,
    save_character as core_save_character,
    load_character as core_load_character
)
from modules.integrations import (
    refine_master,
    generate_name_master,
    generate_image_master,
    get_ollama_models,
    check_comfy_availability
)
from modules.name_generator import generate_fantasy_name

app = FastAPI(title="Chronicle Portrait Studio API")

# Setup CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], # Since frontend might run on different port
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class PromptRequest(BaseModel):
    character_name: str
    features: List[str]
    randomization: List[bool]
    extra_info: List[str]

    def to_args(self):
        return [self.character_name] + self.features + self.randomization + self.extra_info

class NamingRequest(BaseModel):
    race: str

class RegenerationRequest(BaseModel):
    current_values: List[str]
    checkboxes: List[bool]
    
    def to_args(self):
        return self.current_values + self.checkboxes

class RefinementRequest(BaseModel):
    prompt: str
    backend: str
    ollama_model: Optional[str] = None
    hf_text_model: Optional[str] = None
    hf_text_provider: Optional[str] = None
    oauth_token: Optional[str] = None
    character_name: Optional[str] = None

class ImageGenerationRequest(BaseModel):
    refined_prompt: str
    technical_prompt: str
    aspect_ratio: str
    backend: str
    hf_image_model: Optional[str] = None
    hf_image_provider: Optional[str] = None
    oauth_token: Optional[str] = None
    character_name: Optional[str] = None

@app.get("/api/config")
def get_config():
    """Returns static config info needed by frontend to render dropdowns, etc."""
    from modules.core_logic import features_data, get_example_list
    from modules.integrations import gemini_active, hf_active
    from modules.config import HF_TEXT_MODELS, HF_IMAGE_MODELS
    return {
        "features_data": features_data,
        "feature_sequence": FEATURE_SEQUENCE,
        "sections": SECTIONS,
        "ollama_models": get_ollama_models(),
        "comfy_active": check_comfy_availability(),
        "gemini_active": gemini_active,
        "hf_active": hf_active,
        "is_hf_space": bool(os.environ.get("SPACE_ID")),
        "hf_text_models": HF_TEXT_MODELS,
        "hf_image_models": HF_IMAGE_MODELS,
        "examples": get_example_list()
    }

@app.get("/api/example/{filename}")
def get_example(filename: str):
    from modules.config import EXAMPLES_DIR
    import os
    path = os.path.join(EXAMPLES_DIR, filename)
    if os.path.exists(path):
        return FileResponse(path, media_type="application/json")
    raise HTTPException(status_code=404, detail="Example not found")

@app.post("/api/generate_prompt")
def generate_prompt(req: PromptRequest):
    args = req.to_args()
    prompt = core_generate_prompt(*args)
    return {"prompt": prompt}

@app.post("/api/regenerate_features")
def regenerate_features(req: RegenerationRequest):
    args = req.to_args()
    new_values = core_handle_regeneration(*args)
    return {"new_values": new_values}

@app.post("/api/generate_name")
def proxy_generate_name(req: NamingRequest):
    name = generate_fantasy_name(req.race)
    return {"name": name}

@app.post("/api/refine_prompt")
def refine_prompt(req: RefinementRequest):
    result, error_msg = refine_master(
        req.prompt, 
        req.backend, 
        req.ollama_model, 
        req.hf_text_model, 
        req.hf_text_provider, 
        req.oauth_token, 
        req.character_name
    )
    # result can be a gr.update() if it fails, and error_msg will have content
    if error_msg:
        raise HTTPException(status_code=400, detail=error_msg)
    return {"refined_prompt": result}

@app.post("/api/generate_image")
def generate_image(req: ImageGenerationRequest):
    print(f"DEBUG: Received Image Request: {req.dict()}")
    img, img_path, status_msg = generate_image_master(
        req.refined_prompt,
        req.technical_prompt,
        req.aspect_ratio,
        req.backend,
        req.hf_image_model,
        req.hf_image_provider,
        req.oauth_token,
        req.character_name
    )
    if img_path and os.path.exists(img_path):
        return FileResponse(img_path, media_type="image/png", headers={"X-Status-Msg": status_msg})
    else:
        raise HTTPException(status_code=500, detail=status_msg or "Failed to generate image.")

@app.get("/api/config/oauth")
def get_oauth_config():
    return {
        "oauth_client_id": os.environ.get("OAUTH_CLIENT_ID")
    }

@app.post("/api/save_character")
def save_character(req: PromptRequest):
    args = req.to_args()
    file_path = core_save_character(*args)
    if file_path and os.path.exists(file_path):
         return FileResponse(file_path, media_type="application/json", filename=os.path.basename(file_path))
    raise HTTPException(status_code=500, detail="Failed to save character state.")

from fastapi.staticfiles import StaticFiles
frontend_path = os.path.join(os.path.dirname(__file__), "frontend", "out")
if os.path.exists(frontend_path):
    app.mount("/", StaticFiles(directory=frontend_path, html=True), name="frontend")
else:
    print(f"WARNING: Static frontend directory not found at {frontend_path}. Please build the Next.js frontend.")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)