# SPDX-License-Identifier: MIT
"""
AgentCore entrypoint for LlamaIndex Workflows.
Auto-generated by LlamaCloud — wraps workflows for AWS Bedrock AgentCore Runtime.

Uses BedrockAgentCoreApp with @app.entrypoint decorator for native AgentCore integration.
"""

from __future__ import annotations

import functools
import logging
from pathlib import Path
from typing import Any, Literal, cast

from bedrock_agentcore import BedrockAgentCoreApp
from llama_deploy.appserver.workflow_loader import load_workflows
from llama_deploy.core.deployment_config import (
    read_deployment_config_from_git_root_or_cwd,
)
from pydantic import BaseModel, ValidationError
from workflows import Workflow
from workflows.context.serializers import JsonSerializer
from workflows.events import StartEvent

logger = logging.getLogger(__name__)
app = BedrockAgentCoreApp()


class WorkflowResult(BaseModel):
    workflow: str
    status: Literal["completed", "failed"]
    result: str | None = None
    error: str | None = None
    session_id: str | None = None


@functools.lru_cache(maxsize=1)
def _load_workflows() -> tuple[dict[str, Workflow], str, str | None]:
    config_dir = Path.cwd()
    if not (config_dir / "pyproject.toml").exists():
        raise FileNotFoundError(
            "No pyproject.toml found at "
            f"{config_dir}.\n"
            "Add a pyproject.toml to your project and re-run."
        )
    config = read_deployment_config_from_git_root_or_cwd(
        Path.cwd(), config_dir
    )  # let errors bubble up if misconfigured
    workflows = load_workflows(config)
    has_default = any(key == "default" for key in list(workflows.keys()))
    default_workflow = "default" if has_default else next(iter(workflows))
    file_workflow = None
    for name, wf in workflows.items():
        if wf.start_event_class.__name__ == "FileEvent":
            file_workflow = name
            break
    return workflows, default_workflow, file_workflow


async def _run_workflow(
    workflow_name: str, wf: Workflow, start_event: StartEvent
) -> WorkflowResult:
    """Run a workflow by name with the given event data."""
    try:
        handler = wf.run(start_event=start_event)
        result = await handler
    except Exception as e:
        logger.error("Workflow '%s' failed: %s", workflow_name, e, exc_info=True)
        return WorkflowResult(
            workflow=workflow_name, error=f"Workflow failed: {e}", status="failed"
        )

    return WorkflowResult(
        workflow=workflow_name,
        status="completed",
        result=JsonSerializer().serialize(result),
    )


def _parse_and_validate_payload(
    workflows: dict[str, Any],
    default_workflow: str,
    file_workflow: str | None,
    payload: dict,
) -> tuple[str, Workflow, StartEvent] | tuple[str, str]:
    """Parse incoming payload to determine workflow and event data.

    Supports:
        - Explicit: {"workflow": "process-file", "start_event": {"file_id": "123"}}
        - Shorthand: {"file_id": "123"} -> routes to process-file workflow
        - Default: {} -> routes to default workflow
    """
    workflow_name = payload.get("workflow")
    event_data = payload.get("start_event", {})

    if not workflow_name:
        if "file_id" in payload and file_workflow is not None:
            workflow_name = file_workflow
            event_data = {"file_id": payload["file_id"]}
        else:
            workflow_name = default_workflow
    if workflow_name not in workflows:
        return workflow_name, f"Workflow not found: {workflow_name}"
    wf = workflows[workflow_name]
    start_cls = wf.start_event_class
    try:
        data = start_cls.model_validate(event_data)
    except ValidationError as e:
        return workflow_name, f"Invalid input data: {e}"
    return workflow_name, wf, data


@app.entrypoint
async def invoke(payload: dict, context: Any) -> dict[str, Any]:
    """Main AgentCore entrypoint — routes to appropriate workflow."""
    workflows, default_workflow, file_workflow = _load_workflows()
    parsed = _parse_and_validate_payload(
        workflows, default_workflow, file_workflow, payload
    )
    if len(parsed) == 2 and all(isinstance(p, str) for p in parsed):
        workflow_name, error = cast(tuple[str, str], parsed)
        return WorkflowResult(
            workflow=workflow_name,
            error=error,
            status="failed",
            session_id=getattr(
                context, "session_id"
            ),  # stay on the safe side with the session_id attribute on an unknown type
        ).model_dump()
    workflow_name, wf, event = cast(tuple[str, Workflow, StartEvent], parsed)
    result = await _run_workflow(workflow_name, wf, event)
    result.session_id = getattr(context, "session_id")
    return result.model_dump()
