app.py

"""
═══════════════════════════════════════════════════════════
  CORA — Unified Entry Point for Hugging Face Spaces
  Merges the FastAPI backend + Gradio UI into one process.
═══════════════════════════════════════════════════════════
"""

import os
import io
import uuid
import threading

from dotenv import load_dotenv
load_dotenv()

import gradio as gr
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from PIL import Image

from cora_engine import CoraEngine
from cora_curator import CoraCurator
from cora_vision import CoraVision
from cora_memory import CoraMemory

# ━━━ Initialize Core (Lazy) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━
engine = None
curator = None
vision = None
memory = None

from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app: FastAPI):
    global engine, curator, vision, memory
    print("🚀 Initializing Cora Engines...")
    engine = CoraEngine()
    curator = CoraCurator()
    vision = CoraVision()
    memory = CoraMemory()
    yield
    print("💤 Shutting down...")

# ━━━ FastAPI Backend ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
api = FastAPI(
    title="Cora API", 
    description="Historical Archive Generator",
    lifespan=lifespan
)

ARCHIVE_DIR = os.environ.get("ARCHIVE_DIR", "archive_images")
os.makedirs(ARCHIVE_DIR, exist_ok=True)


class AgentPrompt(BaseModel):
    prompt: str
    use_curator: bool = True


class SearchQuery(BaseModel):
    query: str
    limit: int = 10


@api.get("/health")
def health_check():
    """Service status with dynamic backend reporting."""
    return {
        "status": "online",
        "engine": {
            "status": engine.client is not None,
            "backend": "cloud (huggingface)" if engine.client else "none",
            "model": engine.MODEL_ID
        },
        "curator": {
            "status": curator.backend != "none",
            "backend": curator.backend,
            "model": curator.MODEL_ID
        },
        "vision": "online" if vision.clip_model else "offline",
        "memory": "online" if memory.client else "offline"
    }


def archive_generation(image, prompt):
    """Background: save + embed + tag."""
    try:
        filename = f"{uuid.uuid4()}.png"
        filepath = os.path.join(ARCHIVE_DIR, filename)
        image.save(filepath)
        embedding = vision.embed_image(image)
        tags = vision.detect_tags(image)
        memory.save(filepath, embedding, prompt, tags)
        print(f"✅ Archived: {filepath} | tags: {tags}")
    except Exception as e:
        print(f"❌ Archive failed: {e}")


@api.post("/agent/generate")
async def agent_generate(request: AgentPrompt, background_tasks: BackgroundTasks):
    try:
        if not request.prompt or not request.prompt.strip():
            raise HTTPException(status_code=400, detail="Prompt cannot be empty.")

        final_prompt = request.prompt
        if request.use_curator:
            try:
                final_prompt = curator.refine_prompt(request.prompt)
            except Exception:
                final_prompt = request.prompt

        result = engine.generate_from_text(final_prompt)
        
        if "transparent" in request.prompt.lower():
            try:
                from rembg import remove
                result = remove(result)
            except Exception as e:
                print(f"Failed to remove background: {e}")

        background_tasks.add_task(archive_generation, result, final_prompt)

        img_bytes = io.BytesIO()
        result.save(img_bytes, format="PNG")
        from fastapi.responses import Response
        return Response(content=img_bytes.getvalue(), media_type="image/png")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@api.post("/curator/search")
async def curator_search(request: SearchQuery):
    try:
        emb = vision.embed_text(request.query)
        if not emb:
            return {"results": []}

        query_lower = request.query.lower()
        tag_hints = []
        cultural_markers = {
            "roman": ["roman", "rome"],
            "greek": ["greek", "greece", "hellenic"],
            "egyptian": ["egypt", "egyptian"],
            "medieval": ["medieval", "middle ages"],
            "renaissance": ["renaissance"],
        }
        for culture, keywords in cultural_markers.items():
            if any(kw in query_lower for kw in keywords):
                tag_hints.extend(keywords)

        if tag_hints:
            results = memory.search_hybrid(emb, k=request.limit, tag_filter=tag_hints)
        else:
            results = memory.search_by_vector(emb, k=request.limit)

        images = []
        if results["ids"]:
            for i, uid in enumerate(results["ids"][0]):
                meta = results["metadatas"][0][i]
                path = meta.get("path")
                if path and os.path.exists(path):
                    filename = os.path.basename(path)
                    images.append({
                        "path": f"/archive_images/{filename}",
                        "tags": meta.get("tags"),
                        "prompt": meta.get("prompt"),
                        "score": float(results["distances"][0][i]),
                    })
        return {"results": images}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


api.mount("/archive_images", StaticFiles(directory=ARCHIVE_DIR), name="archive_images")

# ━━━ Gradio UI ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

# Use the in-process API directly instead of HTTP calls
API_URL = None  # We call functions directly now


def ui_generate(user_prompt, is_monochromy, monochromy_color, is_transparent, progress=gr.Progress()):
    """Generate an archive illustration — calls engine directly."""
    try:
        if not user_prompt or not user_prompt.strip():
            raise gr.Error("Please enter a subject to generate an archive.")

        progress(0.1, desc="🔍 Curator is analyzing your request...")

        augmented_prompt = user_prompt.strip()
        if is_monochromy:
            augmented_prompt += f", monochromy in {monochromy_color} color"
        if is_transparent:
            augmented_prompt += ", transparent background"

        final_prompt = augmented_prompt
        try:
            final_prompt = curator.refine_prompt(augmented_prompt)
        except Exception:
            pass

        progress(0.3, desc="🎨 Generating historical archive... (1-2 min)")
        result = engine.generate_from_text(final_prompt)

        if "transparent" in augmented_prompt.lower():
            try:
                progress(0.8, desc="✂️ Removing background...")
                from rembg import remove
                result = remove(result)
            except Exception as e:
                print(f"Failed to remove background: {e}")

        # Archive in background
        threading.Thread(
            target=archive_generation, args=(result, final_prompt), daemon=True
        ).start()

        progress(1.0, desc="✅ Generated! Indexing in background...")
        return result
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")


def ui_search(query):
    """Semantic search — calls memory directly."""
    try:
        if not query or not query.strip():
            raise gr.Error("Please enter search keywords.")

        emb = vision.embed_text(query)
        if not emb:
            return []

        results = memory.search_by_vector(emb, k=9)
        gallery_items = []
        if results["ids"]:
            for i, uid in enumerate(results["ids"][0]):
                meta = results["metadatas"][0][i]
                path = meta.get("path")
                if path and os.path.exists(path):
                    gallery_items.append((path, meta.get("tags", "")))

        if not gallery_items:
            gr.Info("🔍 No artifacts found. Try different keywords.")
        return gallery_items
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Search error: {str(e)}")


# ━━━ Build Gradio Interface ━━━━━━━━━━━━━━━━━━━━━━━━━━━

custom_css = """
button { transition: all 0.2s ease-in-out !important; }
button:hover { transform: scale(1.02); opacity: 0.95; }
button:active { transform: scale(0.98); }
button:focus-visible, textarea:focus-visible, input:focus-visible {
    outline: 2px solid #8B5CF6 !important;
    outline-offset: 2px !important;
}
@media (max-width: 768px) { .gallery { gap: 0.5rem !important; } }
"""

with gr.Blocks(
    title="Cora — Historical Archive Generator",
    theme=gr.themes.Monochrome(),
    css=custom_css,
) as demo:
    gr.Markdown(
        f"""
        # 🏛️ Cora — The Archive Curator

        **Generate AI-powered historical illustrations** using SDXL and Visual RAG.
        Transform simple prompts into detailed historical scenes, artifacts, and engravings.

        **How it works**: Your prompt is enhanced by an AI curator, then rendered as a
        historically-themed illustration and archived for future semantic search.
        
        *Active Curator Backend:* **{{curator.backend}}** | *Model:* **{{curator.MODEL_ID}}**
        """
    )

    with gr.Tabs():
        with gr.Tab("🎨 Generate"):
            with gr.Row():
                with gr.Column(scale=1):
                    txt_input = gr.Textbox(
                        label="Historical Subject",
                        placeholder="Describe a historical scene, artifact, or character...",
                        lines=3,
                        info="Be specific! Mention time periods, cultures, contexts, or ask for a monochromy.",
                    )
                    with gr.Row():
                        chk_monochromy = gr.Checkbox(label="Monochromy", value=False)
                        color_picker = gr.ColorPicker(label="Monochromy Color", value="#000000")
                        chk_transparent = gr.Checkbox(label="Transparent Background", value=False)

                    btn_generate = gr.Button(
                        "🔮 Generate Archive", variant="primary", size="lg"
                    )
                    gr.Examples(
                        examples=[
                            ["A woman mending fishing nets by the shore of a medieval village"],
                            ["Baking bread in a communal stone oven, soot and flour on hands"],
                            ["An apprentice carving a wooden chair in a 17th-century workshop"],
                            ["A Roman merchant counting copper coins at a busy market stall"],
                            ["Tending to a small vegetable garden in the courtyard of a Tudor cottage"],
                            ["Spinning wool into yarn using a simple wooden spindle by a hearth"],
                        ],
                        inputs=txt_input,
                        label="💡 Daily Life Examples (click to use)",
                    )
                with gr.Column(scale=1):
                    output_image = gr.Image(
                        label="🖼️ Recovered Artifact", show_label=True
                    )

            btn_generate.click(
                fn=ui_generate,
                inputs=[txt_input, chk_monochromy, color_picker, chk_transparent],
                outputs=output_image,
                show_progress="full",
            )

        with gr.Tab("📚 The Archive"):
            gr.Markdown(
                """
                ### Search the Visual Memory
                Use semantic search to find archived artifacts by keywords, themes, or cultural markers.
                """
            )
            with gr.Row():
                search_box = gr.Textbox(
                    label="Search Query",
                    placeholder="Try: 'roman armor', 'medieval', 'warrior'",
                    scale=4,
                )
                btn_search = gr.Button("🔍 Search", variant="secondary", scale=1)

            gallery = gr.Gallery(
                label="Found Artifacts",
                columns=[1, 2, 3],
                height="auto",
                object_fit="contain",
                show_label=True,
            )
            btn_search.click(fn=ui_search, inputs=[search_box], outputs=gallery)

            gr.Examples(
                examples=[
                    ["roman gladiator"],
                    ["medieval castle"],
                    ["viking warrior"],
                    ["renaissance art"],
                ],
                inputs=search_box,
                label="💡 Quick Search Terms",
            )

# ━━━ Mount FastAPI into Gradio ━━━━━━━━━━━━━━━━━━━━━━━━
app = gr.mount_gradio_app(api, demo, path="/")

# ━━━ Launch ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)