# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_prism

from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, Dict, List, Literal, Optional

from fastapi import FastAPI, HTTPException
from loguru import logger
from pydantic import BaseModel

from coreason_identity.models import UserContext
from coreason_prism.analyst import Analyst
from coreason_prism.biologist import Biologist
from coreason_prism.chemist import Chemist


class AgentTask(BaseModel):
    role: Literal["chemist", "biologist", "analyst"]
    instruction: str
    context: Dict[str, Any]


class AgentResponse(BaseModel):
    status: str
    content: str
    artifacts: Optional[List[Dict[str, Any]]] = None


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    # Startup: Load heavy models
    logger.info("Initializing Agent Microservice... Loading models.")
    try:
        # Load Biologist model (MedSAM)
        logger.info("Pre-loading Biologist models...")
        Biologist._load_model()

        # Load Analyst model (DePlot)
        logger.info("Pre-loading Analyst models...")
        Analyst._load_model()

        # Store in state if needed, but the classes handle their own state (singleton/class attributes)
        # app.state.models could be used if we returned the model instances,
        # but existing pattern uses class methods/attributes.
        # We can just flag them as loaded.
        app.state.models_loaded = True
        logger.info("Models loaded successfully.")
    except Exception as e:
        logger.error(f"Failed to load models during startup: {e}")
        # We might not want to crash startup if one model fails?
        # But prompt says "load them once in the lifespan".
        # If it fails, the service might be unhealthy.
        # For now, we log and continue, but request processing might fail.

    yield

    # Shutdown
    logger.info("Shutting down Agent Microservice.")


app = FastAPI(title="CoReason Prism Agent Service", lifespan=lifespan)


@app.post("/agent/run", response_model=AgentResponse)  # type: ignore[misc]
async def run_agent(task: AgentTask, context: UserContext) -> Any:
    logger.info(f"Received task for role: {task.role}")

    try:
        agent: Any = None
        if task.role == "chemist":
            agent = Chemist()
        elif task.role == "biologist":
            agent = Biologist()
        elif task.role == "analyst":
            agent = Analyst()

        # 'else' is unreachable due to Pydantic validation on AgentTask.role Literal

        result = await agent.run(task.instruction, task.context, user_context=context)
        return result

    except ValueError as ve:
        logger.error(f"Validation error: {ve}")
        raise HTTPException(status_code=400, detail=str(ve)) from ve
    except Exception as e:
        logger.exception("Internal Server Error during agent execution")
        # Return 500 with clean error message
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}") from e
