"""
Meta AI FastAPI Service (v2.0)

Uses metaai-api library for Meta AI image generation.
See: https://github.com/mir-ashiq/metaai-api
"""
from contextlib import asynccontextmanager
from fastapi import FastAPI, BackgroundTasks, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import asyncio
import uuid



from .models import (
    GenerateRequest,
    GenerateResponse,
    ImageResult,
    TaskStatusResponse,
    HealthResponse,
    GrokChatRequest,
    GrokChatResponse
)
from .grok_client import GrokChatClient
from .meta_crawler import meta_crawler

# Initialize Grok client
grok_client = GrokChatClient()


# Task storage (in-memory for simplicity)
tasks: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup and shutdown events"""
    print("[MetaAI] Starting Meta AI service...")
    yield
    print("[MetaAI] Shutting down...")


app = FastAPI(
    title="Meta AI Image Generation Service",
    description="FastAPI wrapper for Meta AI image generation using metaai-api",
    version="2.0.0",
    lifespan=lifespan
)

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


@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    return HealthResponse(
        status="healthy",
        version="2.0.0",
        browser_ready=True  # metaai-api handles this internally
    )


@app.get("/rate-limit")
async def get_rate_limit():
    """Get current rate limiting status"""
    return meta_crawler.get_rate_limit_status()


@app.post("/generate/sync", response_model=GenerateResponse)
async def generate_sync(request: GenerateRequest):
    """
    Synchronous image generation - returns when complete.
    
    Requires:
    - prompt: The image generation prompt
    - cookies: Facebook/Meta cookies (JSON array or string format)
    """
    try:
        images = await meta_crawler.generate_images(
            prompt=request.prompt,
            cookies=request.cookies,
            num_images=request.num_images
        )
        
        return GenerateResponse(
            success=True,
            images=images,
            error=None
        )
        
    except Exception as e:
        return GenerateResponse(
            success=False,
            images=[],
            error=str(e)
        )


@app.post("/generate", response_model=GenerateResponse)
async def generate_async(request: GenerateRequest, background_tasks: BackgroundTasks):
    """
    Async image generation - returns immediately with task_id.
    Poll /status/{task_id} for results.
    """
    task_id = str(uuid.uuid4())
    
    tasks[task_id] = {
        "status": "pending",
        "images": [],
        "error": None
    }
    
    async def run_generation():
        try:
            images = await meta_crawler.generate_images(
                prompt=request.prompt,
                cookies=request.cookies,
                num_images=request.num_images
            )
            tasks[task_id] = {
                "status": "completed",
                "images": images,
                "error": None
            }
        except Exception as e:
            tasks[task_id] = {
                "status": "failed",
                "images": [],
                "error": str(e)
            }
    
    # Run in background
    background_tasks.add_task(asyncio.create_task, run_generation())
    
    return GenerateResponse(
        success=True,
        images=[],
        error=None,
        task_id=task_id
    )


@app.get("/status/{task_id}", response_model=TaskStatusResponse)
async def get_task_status(task_id: str):
    """Get status of async generation task"""
    if task_id not in tasks:
        raise HTTPException(status_code=404, detail="Task not found")
    
    task = tasks[task_id]
    return TaskStatusResponse(
        task_id=task_id,
        status=task["status"],
        images=task["images"],
        error=task["error"]
    )


@app.delete("/status/{task_id}")
async def delete_task(task_id: str):
    """Clean up completed task"""
    if task_id in tasks:
        del tasks[task_id]
        return {"deleted": True}
    raise HTTPException(status_code=404, detail="Task not found")



@app.post("/grok/chat", response_model=GrokChatResponse)
async def grok_chat(request: GrokChatRequest):
    """
    Chat with Grok AI
    """
    try:
        response = await grok_client.chat(request.message, request.history, request.cookies, request.user_agent)
        return GrokChatResponse(response=response)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))