import json
import logging
from contextlib import asynccontextmanager
from typing import List, Dict, Optional, Any
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel

from alayaflow.api import Flow

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("alayaflow.server")

# =======================
# Pydantic 请求模型定义
# =======================

class InitConfig(BaseModel):
    config: Dict[str, Any] = {}

class LoadWorkflowRequest(BaseModel):
    workflow_id: str
    version: str
    init_args: Optional[Dict] = {}
    storage_path: Optional[str] = None

class InstallWorkflowRequest(BaseModel):
    workflow_id: str

class ExecuteWorkflowRequest(BaseModel):
    workflow_id: str
    version: str
    inputs: Dict[str, Any]
    context: Dict[str, Any] = {}
    executor_type: str = "naive"
    metadata: Dict[str, Any] = {}

class ModelProfileRequest(BaseModel):
    profiles: List[Dict[str, Any]]

# =======================
# FastAPI App Setup
# =======================

flow = Flow()

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Initializing AlayaFlow API Singleton...")
    try:
        flow.init()     # all set in .env
        logger.info("Initialization successful.")
    except Exception as e:
        logger.error(f"Initialization failed: {e}")
        raise e
    
    yield
    
    logger.info("Shutting down AlayaFlow API Singleton...")

app = FastAPI(title="AlayaFlow Server", lifespan=lifespan)

# =======================
# API Endpoints
# =======================

@app.get("/health")
async def health_check():
    return {"status": "ok", "initialized": flow.is_inited()}

@app.get("/workflows/available")
async def list_available():
    return flow.list_available_workflows()

@app.get("/workflows/loaded")
async def list_loaded():
    return flow.list_loaded_workflow_ids()

@app.post("/workflows/install")
async def install_workflow(request: InstallWorkflowRequest):
    try:
        success = flow.install_workflow(request.workflow_id)
        if not success:
            raise HTTPException(status_code=400, detail="Installation failed")
        return {"message": "Workflow installed successfully", "workflow_id": request.workflow_id}
    except Exception as e:
        logger.error(f"Error installing workflow: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/workflows/load")
async def load_workflow(request: LoadWorkflowRequest):
    try:
        kwargs = {"init_args": request.init_args}
        if request.storage_path:
            kwargs["storage_path"] = request.storage_path
            
        flow.load_workflow(request.workflow_id, request.version, **kwargs)
        return {"message": "Workflow loaded successfully"}
    except Exception as e:
        logger.error(f"Error loading workflow: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/models/register")
async def register_models(request: ModelProfileRequest):
    try:
        flow.register_models(request.profiles)
        return {"message": "Models registered successfully"}
    except Exception as e:
        logger.error(f"Error registering models: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/workflows/execute/stream")
async def execute_workflow_stream(request: ExecuteWorkflowRequest):
    """
    流式执行工作流 (NDJSON 风格)
    """
    async def event_generator():
        try:
            async for event in flow.exec_workflow_async(
                workflow_id=request.workflow_id,
                version=request.version,
                inputs=request.inputs,
                context=request.context,
                executor_type=request.executor_type,
                metadata=request.metadata
            ):
                yield json.dumps(event, ensure_ascii=False) + "\n"
        except Exception as e:
            logger.error(f"Error during streaming execution: {e}")
            error_event = {"event": "error", "error": str(e)}
            yield json.dumps(error_event, ensure_ascii=False) + "\n"

    return StreamingResponse(event_generator(), media_type="application/x-ndjson")

@app.post("/workflows/execute/sync")
async def execute_workflow_sync(request: ExecuteWorkflowRequest):
    results = []
    try:
        async for event in flow.exec_workflow_async(
            workflow_id=request.workflow_id,
            version=request.version,
            inputs=request.inputs,
            context=request.context,
            executor_type=request.executor_type,
            metadata=request.metadata
        ):
            results.append(event)
        return JSONResponse(content=results)
    except Exception as e:
        logger.error(f"Error during sync execution: {e}")
        raise HTTPException(status_code=500, detail=str(e))

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