api.py

from fastapi import FastAPI, Form, HTTPException, BackgroundTasks
from fastapi.responses import Response
from cora_engine import CoraEngine
from cora_curator import CoraCurator
from cora_vision import CoraVision
from cora_memory import CoraMemory
import io
import os
import uuid

from pydantic import BaseModel

app = FastAPI(title="Cora API", description="Fake Historical Archive Generator")
engine = CoraEngine()
curator = CoraCurator()
vision = CoraVision()
memory = CoraMemory()

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

@app.get("/health")
def health_check():
    """Checks if the engine and HF connection are ready."""
    status = {"status": "online", "model": engine.MODEL_ID}
    if not engine.client:
        status["status"] = "offline (engine)"
    if not curator.client:
        status["curator"] = "offline"
    else:
        status["curator"] = curator.MODEL_ID
    
    # Check Vision/Memory (simple check if initialized)
    status["vision"] = "online" if vision.clip_model else "offline"
    status["memory"] = "online" if memory.client else "offline"
    
    return status

def archive_generation(image, prompt):
    """Helper to save image and metadata to Visual Memory."""
    try:
        filename = f"{uuid.uuid4()}.png"
        filepath = os.path.join("archive_images", filename)
        
        # Save to disk
        image.save(filepath)
        
        # Analyze (Vision)
        embedding = vision.embed_image(image)
        tags = vision.detect_tags(image)
        
        # Save to Memory (Vector DB)
        memory.save(filepath, embedding, prompt, tags)
        print(f"✅ Background Archiving Complete: {filepath} with tags {tags}")
    except Exception as e:
        print(f"❌ Background Archiving Failed: {e}")

@app.post("/agent/generate")
async def agent_generate(request: AgentPrompt, background_tasks: BackgroundTasks):
    """
    Agent-friendly endpoint receiving JSON.
    Returns the raw PNG image.
    """
    try:
        # Validate input
        if not request.prompt or not request.prompt.strip():
            raise HTTPException(
                status_code=400, 
                detail="Prompt cannot be empty. Please provide a description."
            )
        
        # 1. Curate (Refine Prompt)
        final_prompt = request.prompt
        if request.use_curator:
            try:
                final_prompt = curator.refine_prompt(request.prompt)
            except Exception as curator_error:
                print(f"Curator failed: {curator_error}, using original prompt")
                # Fallback to original if curator fails
                final_prompt = request.prompt
        
        # 2. Generate
        result = engine.generate_from_text(final_prompt)
        
        # 3. Archive (Background Task)
        # We pass a copy or the object itself. Since PIL images are in memory,
        # we need to be careful. However, 'result' is a PIL Image.
        # It's safer to pass the image object. background_tasks will run after return.
        background_tasks.add_task(archive_generation, result, final_prompt)
        
        # Return as PNG
        img_byte_arr = io.BytesIO()
        result.save(img_byte_arr, format='PNG')
        return Response(content=img_byte_arr.getvalue(), media_type="image/png")
        
    except HTTPException:
        raise
    except ValueError as e:
        # User input errors
        raise HTTPException(
            status_code=400, 
            detail=f"Invalid request: {str(e)}"
        )
    except RuntimeError as e:
        # Server/API errors
        error_msg = str(e).lower()
        if "timeout" in error_msg or "took too long" in error_msg:
            raise HTTPException(
                status_code=500,
                detail="Image generation timed out. Try a simpler prompt."
            )
        else:
            raise HTTPException(
                status_code=500, 
                detail=f"Generation failed: {str(e)}"
            )
    except Exception as e:
        print(f"Unexpected server error: {e}")
        raise HTTPException(
            status_code=500, 
            detail="An unexpected error occurred. Please try again."
        )

@app.post("/v1/archive")
async def generate_archive(
    background_tasks: BackgroundTasks,
    prompt: str = Form(...)
):
    """
    Generates an 'archive' style image from text.
    """
    try:
        # 1. Curate (Auto-refine for UI)
        enhanced_prompt = curator.refine_prompt(prompt)

        # 2. Generate
        result = engine.generate_from_text(enhanced_prompt)
        
        # 3. Archive (Background Task)
        background_tasks.add_task(archive_generation, result, enhanced_prompt)
        
        # Return as PNG
        img_byte_arr = io.BytesIO()
        result.save(img_byte_arr, format='PNG')
        return Response(content=img_byte_arr.getvalue(), media_type="image/png")
        
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except RuntimeError as e:
        raise HTTPException(status_code=500, detail=str(e))
    except Exception as e:
        print(f"Server Error: {e}")
        raise HTTPException(status_code=500, detail="Internal Server Error")

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

@app.post("/curator/search")
async def curator_search(request: SearchQuery):
    """
    Semantic search for the UI gallery with intelligent filtering.
    """
    try:
        # 1. Embed query
        emb = vision.embed_text(request.query)
        if not emb:
            return {"results": []}
        
        # 2. Extract potential tags from query for filtering
        query_lower = request.query.lower()
        tag_hints = []
        source_hint = None
        
        # Detect cultural/temporal keywords
        cultural_markers = {
            "roman": ["roman", "rome"],
            "greek": ["greek", "greece", "hellenic"],
            "egyptian": ["egypt", "egyptian"],
            "medieval": ["medieval", "middle ages"],
            "renaissance": ["renaissance"], 
            "enlightment century": ["enlightment century"],
            "industrial revolution":["industrial revolution"],
            "modern times" : ["modern times", "20th century", "21st century"],
        }
        
        for culture, keywords in cultural_markers.items():
            if any(kw in query_lower for kw in keywords):
                tag_hints.extend(keywords)
        
        # 3. Use hybrid search if we detected cultural markers
        if tag_hints:
            results = memory.search_hybrid(emb, k=request.limit, tag_filter=tag_hints)
        else:
            # Fallback to pure semantic if no specific markers
            results = memory.search_by_vector(emb, k=request.limit)
        
        # 4. Format result
        images = []
        if results['ids']:
            ids = results['ids'][0]
            metadatas = results['metadatas'][0]
            distances = results['distances'][0]
            
            for i, uid in enumerate(ids):
                path = metadatas[i].get('path')
                tags = metadatas[i].get('tags')
                prompt = metadatas[i].get('prompt')
                if path and os.path.exists(path):
                    # Convert local path to URL
                    filename = os.path.basename(path)
                    image_url = f"http://localhost:8000/archive_images/{filename}"
                    
                    images.append({
                        "path": image_url,  # Now a URL, not a local path
                        "tags": tags,
                        "prompt": prompt,
                        "score": float(distances[i])
                    })
        return {"results": images}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Mount static files to serve images to UI if needed
from fastapi.staticfiles import StaticFiles
if not os.path.exists("archive_images"):
    os.makedirs("archive_images")
app.mount("/archive_images", StaticFiles(directory="archive_images"), name="archive_images")

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