import json
import logging
from contextlib import asynccontextmanager
from typing import List, Dict, Optional, Any, Generic, TypeVar
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, model_serializer

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] = {}
    workflow_dir: Optional[str] = None

class InstallWorkflowRequest(BaseModel):
    workflow_id: str
    version: 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]]

T = TypeVar("T")

class BaseResponse(BaseModel, Generic[T]):
    @model_serializer(mode='wrap')
    def serialize_model(self, handler) -> dict[str, Any]:
        dumped = handler(self)
        return {k: v for k, v in dumped.items() if v is not None}

    status: str
    data: Optional[T] = None
    error_msg: Optional[str] = None

    @classmethod
    def success(cls, data: Optional[T] = None) -> "BaseResponse[T]":
        return cls(status="success", data=data)

    @classmethod
    def error(cls, error_msg: str) -> "BaseResponse[None]":
        return cls(status="error", error_msg=error_msg)

# =======================
# 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", response_model=BaseResponse[Dict])
async def health_check():
    return BaseResponse.success(data={"initialized": flow.is_inited()})


@app.get("/workflow/list_all", response_model=BaseResponse[Dict])
async def list_all_workflow():
    try:
        workflows = flow.list_all_workflows()
        return BaseResponse.success(data=workflows)
    except Exception as e:
        logger.error(f"Error listing all workflows: {e}", exc_info=True)
        return BaseResponse.error(error_msg=str(e))

@app.get("/workflows/list_loaded", response_model=BaseResponse[List[Dict]])
async def list_loaded_workflow():
    try:
        workflows = flow.list_loaded_workflows()
        return BaseResponse.success(data=workflows)
    except Exception as e:
        logger.error(f"Error listing loaded workflows: {e}", exc_info=True)
        return BaseResponse.error(error_msg=str(e))

@app.get("/workflow/list_local", response_model=BaseResponse[List[Dict]])
async def list_local_workflow():
    try:
        workflows = flow.list_local_workflows()
        return BaseResponse.success(data=workflows)
    except Exception as e:
        logger.error(f"Error listing local workflows: {e}", exc_info=True)
        return BaseResponse.error(error_msg=str(e))

@app.get("/workflow/list_remote", response_model=BaseResponse[List[Dict]])
async def list_remote_workflow():
    try:
        workflows = flow.list_remote_workflows()
        return BaseResponse.success(data=workflows)
    except Exception as e:
        logger.error(f"Error listing remote workflows: {e}", exc_info=True)
        return BaseResponse.error(error_msg=str(e))

@app.get("/workflow/info", response_model=BaseResponse[Dict])
async def get_workflow_info(workflow_id: str, version: str):
    try:
        info = flow.get_workflow_info(workflow_id, version)
        return BaseResponse.success(data=info)
    except Exception as e:
        logger.error(f"Error getting workflow info: {e}", exc_info=True)
        return BaseResponse.error(error_msg=str(e))

@app.post("/workflow/install", response_model=BaseResponse[Dict])
async def install_workflow(request: InstallWorkflowRequest):
    try:
        success = flow.install_workflow(request.workflow_id, request.version)
        if not success:
            raise BaseResponse.error(error_msg="Installation failed")
        return BaseResponse.success(data={
            "message": "Workflow installed successfully",
            "workflow_id": request.workflow_id
        })
    except Exception as e:
        logger.error(f"Error installing workflow: {e}")
        return BaseResponse.error(error_msg=str(e))

@app.post("/workflow/load", response_model=BaseResponse[Dict])
async def load_workflow(request: LoadWorkflowRequest):
    try:
        kwargs = {"init_args": request.init_args}
        if request.workflow_dir:
            kwargs["workflow_dir"] = request.workflow_dir
            
        flow.load_workflow(request.workflow_id, request.version, **kwargs)
        return BaseResponse.success(data={
            "message": "Workflow loaded successfully"
        })
    except Exception as e:
        logger.error(f"Error loading workflow: {e}")
        return BaseResponse.error(error_msg=str(e))

@app.post("/model/register", response_model=BaseResponse[Dict])
async def register_models(request: ModelProfileRequest):
    try:
        flow.register_models(request.profiles)
        return BaseResponse.success(data={"message": "Models registered successfully"})
    except Exception as e:
        logger.error(f"Error registering models: {e}")
        raise BaseResponse.error(error_msg=str(e))

@app.post("/workflow/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", response_model=BaseResponse[Dict])
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 BaseResponse.success(data=results)
    except Exception as e:
        logger.error(f"Error during sync execution: {e}", exc_info=True)
        raise BaseResponse.error(error_msg=str(e))

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