temp_fol_upload/api_server.py

"""
FastAPI backend for Magic Hour AI - streams agent reasoning and results via SSE
LangGraph version: Uses LangChain/LangGraph with fal.ai LLM
"""
import fal_client
from dotenv import load_dotenv
import os
import sys
import json
import asyncio
import threading
import time
import tempfile
import io
import re
import base64
from pathlib import Path
from typing import AsyncGenerator, Optional, List
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from PIL import Image

# Add the mh_langgraph_workflow to path
WORKFLOW_PATH = Path(__file__).parent.parent / "mh_langgraph_workflow"
sys.path.insert(0, str(WORKFLOW_PATH))

# Load environment variables
load_dotenv(WORKFLOW_PATH / ".env")

# Configure fal.ai
fal_client.api_key = os.getenv("FAL_KEY")

# Import from mh_langgraph_workflow
from src.agents import create_agent, ContentAgent
from src.models import FalAILLM

# Global agent instance
agent: Optional[ContentAgent] = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize agent on startup"""
    global agent
    print("🚀 Initializing Magic Hour AI Agent (LangGraph)...")
    try:
        agent = create_agent(
            fal_model_name=os.getenv("FAL_MODEL_NAME", "google/gemini-2.5-flash"),
            temperature=float(os.getenv("LLM_TEMPERATURE", "0.7")),
            max_tokens=int(os.getenv("LLM_MAX_TOKENS", "4096")),
        )
        print(f"✅ Agent initialized with model: {os.getenv('FAL_MODEL_NAME', 'google/gemini-2.5-flash')}")
    except Exception as e:
        print(f"❌ Failed to initialize agent: {e}")
        raise
    yield
    print("👋 Shutting down...")


app = FastAPI(title="Magic Hour AI (LangGraph)", lifespan=lifespan)

# Detect production mode
PRODUCTION_MODE = os.getenv("PRODUCTION", "false").lower() == "true"

# CORS configuration
allowed_origins = ["http://localhost:5173", "http://127.0.0.1:5173", "http://localhost:3000"]
if PRODUCTION_MODE:
    production_domain = os.getenv("FRONTEND_URL", "*")
    allowed_origins = ["*"] if production_domain == "*" else allowed_origins + [production_domain]

app.add_middleware(
    CORSMiddleware,
    allow_origins=allowed_origins,
    allow_credentials=True,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
    expose_headers=["*"],
)

# Serve static frontend in production
if PRODUCTION_MODE:
    frontend_dist = Path(__file__).parent / "dist"
    if frontend_dist.exists():
        app.mount("/assets", StaticFiles(directory=str(frontend_dist / "assets")), name="assets")
        print(f"✅ Serving frontend from {frontend_dist}")


class ChatRequest(BaseModel):
    message: str
    settings: dict = {}
    history: list = []
    session_id: Optional[str] = None  # Optional: Frontend provides to maintain conversation history
    images: Optional[List[str]] = None  # Base64 encoded images from user upload


def save_base64_images(base64_images: List[str]) -> List[str]:
    """Save base64 images to temp files and return paths."""
    saved_paths = []
    for i, b64 in enumerate(base64_images):
        try:
            # Remove data URL prefix if present
            if ',' in b64:
                b64 = b64.split(',')[1]

            image_data = base64.b64decode(b64)
            temp_file = tempfile.NamedTemporaryFile(
                suffix='.png',
                delete=False,
                prefix='user_upload_'
            )
            temp_file.write(image_data)
            temp_file.close()
            saved_paths.append(temp_file.name)
            print(f"DEBUG - Saved uploaded image to: {temp_file.name}")
        except Exception as e:
            print(f"Error saving uploaded image {i}: {e}")
    return saved_paths


class SSEEvent:
    @staticmethod
    def format(event_type: str, data: dict) -> str:
        return f"data: {json.dumps({'type': event_type, **data})}\n\n"


def create_blur_thumbnail(image_path: str, size: int = 20) -> Optional[str]:
    """Create tiny base64 thumbnail for blur effect"""
    try:
        img = Image.open(image_path)
        img.thumbnail((size, size))
        buffer = io.BytesIO()
        img.save(buffer, format="JPEG", quality=50)
        return f"data:image/jpeg;base64,{base64.b64encode(buffer.getvalue()).decode()}"
    except Exception as e:
        print(f"Error creating blur thumbnail: {e}")
        return None


def extract_images_from_history(history: list) -> tuple[List[Image.Image], List[dict], List[str]]:
    """
    Extract PIL images, metadata, and paths from conversation history.
    Returns: (pil_images, image_metadata, image_paths)
    """
    pil_images = []
    image_metadata = []
    image_paths = []
    last_user_prompt = ""

    if not history:
        return [], [], []

    for msg in history:
        if msg.get("role") == "user":
            last_user_prompt = msg.get("content", "")
        elif msg.get("role") == "assistant":
            for img_url in msg.get("images", []):
                if "path=" in img_url:
                    local_path = img_url.split("path=")[-1]
                    try:
                        img = Image.open(local_path)
                        pil_images.append(img.copy())
                        image_paths.append(local_path)
                        image_metadata.append({
                            "index": len(pil_images),
                            "path": local_path,
                            "prompt": last_user_prompt
                        })
                    except Exception as e:
                        print(f"Could not load image {local_path}: {e}")

    return pil_images, image_metadata, image_paths


def parse_image_paths(text: str) -> List[str]:
    """Extract image file paths from text."""
    patterns = [
        r'(/tmp/[^\s\'"]+\.(?:png|jpg|jpeg|webp))',
        r'(/var/folders/[^\s\'"]+\.(?:png|jpg|jpeg|webp))',
    ]
    paths = []
    for pattern in patterns:
        matches = re.findall(pattern, text, re.IGNORECASE)
        paths.extend(matches)
    return [p for p in set(paths) if os.path.exists(p)]


def parse_video_paths(text: str) -> List[str]:
    """Extract video file paths from text."""
    patterns = [
        r'(/tmp/[^\s\'"]+\.(?:mp4|webm|mov))',
        r'(/var/folders/[^\s\'"]+\.(?:mp4|webm|mov))',
    ]
    paths = []
    for pattern in patterns:
        matches = re.findall(pattern, text, re.IGNORECASE)
        paths.extend(matches)
    return [p for p in set(paths) if os.path.exists(p)]


async def generate_with_streaming(
    message: str,
    user_settings: dict,
    history: list = None,
    session_id: str = None,
    uploaded_images: List[str] = None
) -> AsyncGenerator[str, None]:
    """Run the LangGraph agent and stream results via SSE."""
    global agent

    if not agent:
        yield SSEEvent.format("error", {"message": "Agent not initialized"})
        return

    # Use session_id if provided, otherwise create one from timestamp
    # IMPORTANT: Frontend should maintain session_id to preserve conversation history
    if not session_id:
        session_id = f"session-{int(time.time() * 1000)}"
        print(f"DEBUG - Created new session ID: {session_id}")
    else:
        print(f"DEBUG - Using existing session ID: {session_id}")

    mode = user_settings.get("mode", "fast")
    aspect_ratio = user_settings.get("aspectRatio", "square")
    selected_image = user_settings.get("selectedImage")

    # Extract images from history unless the user uploaded new images
    pil_images, image_metadata, image_paths = extract_images_from_history(history)
    if uploaded_images:
        image_paths = list(uploaded_images)
        image_metadata = []
        pil_images = []
        print(f"DEBUG - Using {len(uploaded_images)} uploaded images (ignoring history images)")

    # Just send the user message
    yield SSEEvent.format("reasoning", {"content": "🧠 Processing your request..."})

    full_prompt = f"""{message}"""

    yield SSEEvent.format("reasoning", {"content": "🤖 Starting generation..."})

    generated_images = []
    generated_videos = []
    reasoning_steps = []
    final_response = ""
    last_message_count = 0
    latest_generation_history = []

    loop = asyncio.get_running_loop()
    queue: asyncio.Queue = asyncio.Queue()

    def worker():
        try:
            for event in agent.stream(
                message=full_prompt,
                image_paths=image_paths if image_paths else None,
                thread_id=session_id,
                settings=user_settings,
            ):
                asyncio.run_coroutine_threadsafe(queue.put(event), loop)
        except Exception as e:
            asyncio.run_coroutine_threadsafe(queue.put({"__error__": str(e)}), loop)
        finally:
            asyncio.run_coroutine_threadsafe(queue.put(None), loop)

    threading.Thread(target=worker, daemon=True).start()

    while True:
        event = await queue.get()
        if event is None:
            break
        if isinstance(event, dict) and event.get("__error__"):
            yield SSEEvent.format("error", {"message": event["__error__"]})
            return

        state = event or {}
        if not isinstance(state, dict):
            continue

        generation_history = state.get("generation_history")
        if generation_history is not None:
            latest_generation_history = generation_history

        messages = state.get("messages", [])
        if not isinstance(messages, list):
            continue

        if len(messages) <= last_message_count:
            continue

        new_messages = messages[last_message_count:]
        last_message_count = len(messages)

        for msg in new_messages:
            msg_content = msg.content if hasattr(msg, "content") else str(msg)
            msg_type = type(msg).__name__

            if msg_type == "ReasoningMessage":
                content = str(msg_content)
                reasoning_steps.append(content)
                yield SSEEvent.format("thinking", {"content": content + "\n\n"})
                await asyncio.sleep(0.01)

            elif msg_type == "VisualAnalysisMessage":
                yield SSEEvent.format("visual_analysis", {
                    "content": str(msg_content),
                    "collapsible": True
                })
                await asyncio.sleep(0.01)

            elif msg_type == "AIMessage":
                content_str = str(msg_content)

                # Only extract paths from TOOL RESULT messages (not tool calls)
                # Tool results contain "Successfully" while tool calls contain the original paths
                is_tool_result = "Successfully" in content_str

                if is_tool_result:
                    found_images = parse_image_paths(content_str)
                    found_videos = parse_video_paths(content_str)

                    for img_path in found_images:
                        if img_path not in generated_images:
                            generated_images.append(img_path)
                            yield SSEEvent.format("image_complete", {"url": f"/api/media?path={img_path}"})
                            await asyncio.sleep(0.05)

                    for vid_path in found_videos:
                        if vid_path not in generated_videos:
                            generated_videos.append(vid_path)
                            yield SSEEvent.format("video_complete", {"url": f"/api/media?path={vid_path}"})
                            await asyncio.sleep(0.05)

                # Check for tool calls in response (agent reasoning)
                elif '{"tool":' in content_str or '"tool"' in content_str:
                    try:
                        json_match = re.search(r'\{"tool":\s*"([^"]+)"', content_str)
                        if json_match:
                            yield SSEEvent.format("tool_selected", {
                                "tool": json_match.group(1),
                                "details": ""
                            })
                    except Exception:
                        pass

                    if "[REASONING]" in content_str:
                        reasoning_match = re.search(r'\[REASONING\](.*?)(?=\[ACTION\]|\{|```)', content_str, re.DOTALL)
                        if reasoning_match:
                            reasoning = reasoning_match.group(1).strip()
                            if reasoning and reasoning not in reasoning_steps:
                                reasoning_steps.append(reasoning)
                                yield SSEEvent.format("reasoning_step", {
                                    "content": reasoning,
                                    "collapsible": True
                                })
                    elif "```json" in content_str:
                        reasoning = content_str.split("```json")[0].strip()
                        if reasoning and reasoning not in reasoning_steps:
                            reasoning_steps.append(reasoning)
                            yield SSEEvent.format("reasoning_step", {
                                "content": reasoning[:500],
                                "collapsible": True
                            })
                else:
                    final_response = content_str

    if not generated_images and not generated_videos and latest_generation_history:
        latest_batch = latest_generation_history[-1]
        if latest_batch["type"] == "image":
            for path in latest_batch["paths"]:
                if path not in generated_images and os.path.exists(path):
                    generated_images.append(path)
                    thumbnail = create_blur_thumbnail(path)
                    if thumbnail:
                        yield SSEEvent.format("image_preview", {"blur_data": thumbnail})
                    yield SSEEvent.format("image_complete", {"url": f"/api/media?path={path}"})
        elif latest_batch["type"] == "video":
            for path in latest_batch["paths"]:
                if path not in generated_videos and os.path.exists(path):
                    generated_videos.append(path)
                    yield SSEEvent.format("video_complete", {"url": f"/api/media?path={path}"})

    if final_response:
        clean_response = final_response
        for path in generated_images + generated_videos:
            clean_response = clean_response.replace(path, "").strip()

        if "```json" in clean_response:
            clean_response = clean_response.split("```json")[0].strip()

        if clean_response and len(clean_response) > 5:
            yield SSEEvent.format("agent_message", {"content": clean_response})

    yield SSEEvent.format("reasoning", {"content": "✅ Complete!"})
    yield SSEEvent.format("complete", {
        "session_id": session_id,
        "generated_images": generated_images,
        "generated_videos": generated_videos
    })


@app.post("/api/chat")
async def chat(request: ChatRequest):
    # Save uploaded images to temp files
    uploaded_paths = []
    if request.images:
        uploaded_paths = save_base64_images(request.images)

    return StreamingResponse(
        generate_with_streaming(
            request.message,
            request.settings,
            request.history,
            request.session_id,
            uploaded_paths
        ),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"
        }
    )


@app.get("/api/media")
async def serve_media(path: str):
    """Serve generated media files with security validation"""
    try:
        resolved_path = Path(path).resolve()
        temp_dir = Path(tempfile.gettempdir()).resolve()

        if not str(resolved_path).startswith(str(temp_dir)):
            raise HTTPException(status_code=403, detail="Access denied")

        if not resolved_path.exists():
            raise HTTPException(status_code=404, detail="File not found")

        content_types = {
            ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
            ".webp": "image/webp", ".mp4": "video/mp4", ".webm": "video/webm"
        }
        ext = resolved_path.suffix.lower()
        if ext not in content_types:
            raise HTTPException(status_code=400, detail="Unsupported file type")

        return FileResponse(str(resolved_path), media_type=content_types[ext])
    except HTTPException:
        raise
    except Exception as e:
        print(f"Error serving media: {e}")
        raise HTTPException(status_code=500, detail="Error serving file")


@app.get("/api/health")
async def health():
    return {"status": "ok", "agent_ready": agent is not None, "backend": "langgraph"}


if PRODUCTION_MODE:
    @app.get("/{full_path:path}")
    async def serve_frontend(full_path: str):
        frontend_dist = Path(__file__).parent / "dist"
        if full_path.startswith("api/"):
            raise HTTPException(status_code=404, detail="API endpoint not found")

        requested_file = frontend_dist / full_path
        if requested_file.exists() and requested_file.is_file():
            return FileResponse(requested_file)

        index_file = frontend_dist / "index.html"
        if index_file.exists():
            return FileResponse(index_file)

        raise HTTPException(status_code=404, detail="Frontend not built")


if __name__ == "__main__":
    import uvicorn
    print("Starting Magic Hour AI Server (LangGraph)...")
    print("API: http://localhost:8000")
    print("Frontend: http://localhost:5173")
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")