import logging
from typing import Any, Optional, Union

from fastapi import APIRouter, HTTPException, Path, Query
from pydantic import ValidationError
from sqlalchemy import select
from sqlalchemy.sql import Select
from starlette.requests import Request
from starlette.status import HTTP_404_NOT_FOUND, HTTP_422_UNPROCESSABLE_ENTITY
from strawberry.relay import GlobalID
from typing_extensions import TypeAlias, assert_never

from phoenix.db import models
from phoenix.db.types.identifier import Identifier
from phoenix.db.types.model_provider import ModelProvider
from phoenix.server.api.helpers.prompts.models import (
    PromptInvocationParameters,
    PromptResponseFormat,
    PromptTemplate,
    PromptTemplateFormat,
    PromptTemplateType,
    PromptTools,
)
from phoenix.server.api.routers.v1.models import V1RoutesBaseModel
from phoenix.server.api.routers.v1.utils import ResponseBody, add_errors_to_responses
from phoenix.server.api.types.node import from_global_id_with_expected_type
from phoenix.server.api.types.Prompt import Prompt as PromptNodeType
from phoenix.server.api.types.PromptVersion import PromptVersion as PromptVersionNodeType
from phoenix.server.bearer_auth import PhoenixUser

logger = logging.getLogger(__name__)


class PromptData(V1RoutesBaseModel):
    name: Identifier
    description: Optional[str] = None
    source_prompt_id: Optional[str] = None


class Prompt(PromptData):
    id: str


class PromptVersionData(V1RoutesBaseModel):
    description: Optional[str] = None
    model_provider: ModelProvider
    model_name: str
    template: PromptTemplate
    template_type: PromptTemplateType
    template_format: PromptTemplateFormat
    invocation_parameters: PromptInvocationParameters
    tools: Optional[PromptTools] = None
    response_format: Optional[PromptResponseFormat] = None


class PromptVersion(PromptVersionData):
    id: str


class GetPromptResponseBody(ResponseBody[PromptVersion]):
    pass


class GetPromptsResponseBody(ResponseBody[list[Prompt]]):
    pass


class GetPromptVersionsResponseBody(ResponseBody[list[PromptVersion]]):
    pass


class CreatePromptRequestBody(V1RoutesBaseModel):
    prompt: PromptData
    version: PromptVersionData


class CreatePromptResponseBody(ResponseBody[PromptVersion]):
    pass


router = APIRouter(tags=["prompts"])


@router.get(
    "/prompts",
    operation_id="getPrompts",
    summary="Get all prompts",
    responses=add_errors_to_responses(
        [
            HTTP_422_UNPROCESSABLE_ENTITY,
        ]
    ),
)
async def get_prompts(
    request: Request,
    cursor: Optional[str] = Query(
        default=None,
        description="Cursor for pagination (base64-encoded prompt ID)",
    ),
    limit: int = Query(
        default=100, description="The max number of prompts to return at a time.", gt=0
    ),
) -> GetPromptsResponseBody:
    async with request.app.state.db() as session:
        query = select(models.Prompt).order_by(models.Prompt.id.desc())

        if cursor:
            try:
                cursor_id = GlobalID.from_id(cursor).node_id
                query = query.filter(models.Prompt.id <= int(cursor_id))
            except ValueError:
                raise HTTPException(
                    detail=f"Invalid cursor format: {cursor}",
                    status_code=HTTP_422_UNPROCESSABLE_ENTITY,
                )

        query = query.limit(limit + 1)
        result = await session.execute(query)
        orm_prompts = result.scalars().all()

        if not orm_prompts:
            return GetPromptsResponseBody(next_cursor=None, data=[])

        next_cursor = None
        if len(orm_prompts) == limit + 1:
            last_prompt = orm_prompts[-1]
            next_cursor = str(GlobalID(PromptNodeType.__name__, str(last_prompt.id)))
            orm_prompts = orm_prompts[:-1]

        prompts = [_prompt_from_orm_prompt(orm_prompt) for orm_prompt in orm_prompts]
    return GetPromptsResponseBody(next_cursor=next_cursor, data=prompts)


@router.get(
    "/prompts/{prompt_identifier}/versions",
    operation_id="listPromptVersions",
    summary="List all prompt versions for a given prompt",
    responses=add_errors_to_responses([HTTP_422_UNPROCESSABLE_ENTITY]),
    response_model_by_alias=True,
    response_model_exclude_defaults=True,
    response_model_exclude_unset=True,
)
async def list_prompt_versions(
    request: Request,
    prompt_identifier: str = Path(description="The identifier of the prompt, i.e. name or ID."),
    cursor: Optional[str] = Query(
        default=None,
        description="Cursor for pagination (base64-encoded promptVersion ID)",
    ),
    limit: int = Query(
        default=100, description="The max number of prompt versions to return at a time.", gt=0
    ),
) -> GetPromptVersionsResponseBody:
    query = select(models.PromptVersion)
    query = _filter_by_prompt_identifier(query.join(models.Prompt), prompt_identifier)
    query = query.order_by(models.PromptVersion.id.desc())

    async with request.app.state.db() as session:
        if cursor:
            try:
                cursor_id = GlobalID.from_id(cursor).node_id
                query = query.filter(models.PromptVersion.id <= int(cursor_id))
            except ValueError:
                raise HTTPException(
                    detail=f"Invalid cursor format: {cursor}",
                    status_code=HTTP_422_UNPROCESSABLE_ENTITY,
                )

        query = query.limit(limit + 1)
        result = await session.execute(query)
        orm_versions = result.scalars().all()

        if not orm_versions:
            return GetPromptVersionsResponseBody(next_cursor=None, data=[])

        next_cursor = None
        if len(orm_versions) == limit + 1:
            last_version = orm_versions[-1]
            next_cursor = str(GlobalID(PromptVersionNodeType.__name__, str(last_version.id)))
            orm_versions = orm_versions[:-1]

        versions = [_prompt_version_from_orm_version(orm_version) for orm_version in orm_versions]
        return GetPromptVersionsResponseBody(next_cursor=next_cursor, data=versions)


@router.get(
    "/prompt_versions/{prompt_version_id}",
    operation_id="getPromptVersionByPromptVersionId",
    summary="Get prompt by prompt version ID",
    responses=add_errors_to_responses(
        [
            HTTP_404_NOT_FOUND,
            HTTP_422_UNPROCESSABLE_ENTITY,
        ]
    ),
    response_model_by_alias=True,
    response_model_exclude_defaults=True,
    response_model_exclude_unset=True,
)
async def get_prompt_version_by_prompt_version_id(
    request: Request,
    prompt_version_id: str = Path(description="The ID of the prompt version."),
) -> GetPromptResponseBody:
    try:
        id_ = from_global_id_with_expected_type(
            GlobalID.from_id(prompt_version_id),
            PromptVersionNodeType.__name__,
        )
    except ValueError:
        raise HTTPException(HTTP_422_UNPROCESSABLE_ENTITY, "Invalid prompt version ID")
    async with request.app.state.db() as session:
        prompt_version = await session.get(models.PromptVersion, id_)
        if prompt_version is None:
            raise HTTPException(HTTP_404_NOT_FOUND)
    data = _prompt_version_from_orm_version(prompt_version)
    return GetPromptResponseBody(data=data)


@router.get(
    "/prompts/{prompt_identifier}/tags/{tag_name}",
    operation_id="getPromptVersionByTagName",
    summary="Get prompt by tag name",
    responses=add_errors_to_responses(
        [
            HTTP_404_NOT_FOUND,
            HTTP_422_UNPROCESSABLE_ENTITY,
        ]
    ),
    response_model_by_alias=True,
    response_model_exclude_unset=True,
    response_model_exclude_defaults=True,
)
async def get_prompt_version_by_tag_name(
    request: Request,
    prompt_identifier: str = Path(description="The identifier of the prompt, i.e. name or ID."),
    tag_name: str = Path(description="The tag of the prompt version"),
) -> GetPromptResponseBody:
    try:
        name = Identifier.model_validate(tag_name)
    except ValidationError:
        raise HTTPException(HTTP_422_UNPROCESSABLE_ENTITY, "Invalid tag name")
    stmt = (
        select(models.PromptVersion)
        .join_from(models.PromptVersion, models.PromptVersionTag)
        .where(models.PromptVersionTag.name == name)
    )
    stmt = _filter_by_prompt_identifier(stmt.join(models.Prompt), prompt_identifier)
    async with request.app.state.db() as session:
        prompt_version: models.PromptVersion = await session.scalar(stmt)
        if prompt_version is None:
            raise HTTPException(HTTP_404_NOT_FOUND)
    data = _prompt_version_from_orm_version(prompt_version)
    return GetPromptResponseBody(data=data)


@router.get(
    "/prompts/{prompt_identifier}/latest",
    operation_id="getPromptVersionLatest",
    summary="Get the latest prompt version",
    responses=add_errors_to_responses(
        [
            HTTP_404_NOT_FOUND,
            HTTP_422_UNPROCESSABLE_ENTITY,
        ]
    ),
    response_model_by_alias=True,
    response_model_exclude_defaults=True,
    response_model_exclude_unset=True,
)
async def get_prompt_version_by_latest(
    request: Request,
    prompt_identifier: str = Path(description="The identifier of the prompt, i.e. name or ID."),
) -> GetPromptResponseBody:
    stmt = select(models.PromptVersion).order_by(models.PromptVersion.id.desc()).limit(1)
    stmt = _filter_by_prompt_identifier(stmt.join(models.Prompt), prompt_identifier)
    async with request.app.state.db() as session:
        prompt_version: models.PromptVersion = await session.scalar(stmt)
        if prompt_version is None:
            raise HTTPException(HTTP_404_NOT_FOUND)
    data = _prompt_version_from_orm_version(prompt_version)
    return GetPromptResponseBody(data=data)


@router.post(
    "/prompts",
    operation_id="postPromptVersion",
    summary="Create a prompt version",
    responses=add_errors_to_responses(
        [
            HTTP_422_UNPROCESSABLE_ENTITY,
        ]
    ),
    response_model_by_alias=True,
    response_model_exclude_defaults=True,
    response_model_exclude_unset=True,
)
async def create_prompt(
    request: Request,
    request_body: CreatePromptRequestBody,
) -> CreatePromptResponseBody:
    if (
        request_body.version.template.type.lower() != "chat"
        or request_body.version.template_type != PromptTemplateType.CHAT
    ):
        raise HTTPException(
            HTTP_422_UNPROCESSABLE_ENTITY,
            "Only CHAT template type is supported for prompts",
        )
    prompt = request_body.prompt
    try:
        name = Identifier.model_validate(prompt.name)
    except ValidationError as e:
        raise HTTPException(
            HTTP_422_UNPROCESSABLE_ENTITY,
            "Invalid name identifier for prompt: " + e.errors()[0]["msg"],
        )
    version = request_body.version
    user_id: Optional[int] = None
    if request.app.state.authentication_enabled:
        assert isinstance(user := request.user, PhoenixUser)
        user_id = int(user.identity)
    async with request.app.state.db() as session:
        if not (prompt_id := await session.scalar(select(models.Prompt.id).filter_by(name=name))):
            prompt_orm = models.Prompt(
                name=name,
                description=prompt.description,
            )
            session.add(prompt_orm)
            await session.flush()
            prompt_id = prompt_orm.id
        version_orm = models.PromptVersion(
            user_id=user_id,
            prompt_id=prompt_id,
            description=version.description,
            model_provider=version.model_provider,
            model_name=version.model_name,
            template_type=version.template_type,
            template_format=version.template_format,
            template=version.template,
            invocation_parameters=version.invocation_parameters,
            tools=version.tools,
            response_format=version.response_format,
        )
        session.add(version_orm)
    data = _prompt_version_from_orm_version(version_orm)
    return CreatePromptResponseBody(data=data)


class _PromptId(int): ...


_PromptIdentifier: TypeAlias = Union[_PromptId, Identifier]


def _parse_prompt_identifier(
    prompt_identifier: str,
) -> _PromptIdentifier:
    if not prompt_identifier:
        raise HTTPException(HTTP_422_UNPROCESSABLE_ENTITY, "Invalid prompt identifier")
    try:
        prompt_id = from_global_id_with_expected_type(
            GlobalID.from_id(prompt_identifier),
            PromptNodeType.__name__,
        )
    except ValueError:
        try:
            return Identifier.model_validate(prompt_identifier)
        except ValidationError:
            raise HTTPException(HTTP_422_UNPROCESSABLE_ENTITY, "Invalid prompt name")
    return _PromptId(prompt_id)


def _filter_by_prompt_identifier(
    stmt: Select[tuple[models.PromptVersion]],
    prompt_identifier: str,
) -> Any:
    identifier = _parse_prompt_identifier(prompt_identifier)
    if isinstance(identifier, _PromptId):
        return stmt.where(models.Prompt.id == int(identifier))
    if isinstance(identifier, Identifier):
        return stmt.where(models.Prompt.name == identifier)
    assert_never(identifier)


def _prompt_version_from_orm_version(
    prompt_version: models.PromptVersion,
) -> PromptVersion:
    prompt_template_type = PromptTemplateType(prompt_version.template_type)
    prompt_template_format = PromptTemplateFormat(prompt_version.template_format)
    return PromptVersion(
        id=str(GlobalID(PromptVersionNodeType.__name__, str(prompt_version.id))),
        description=prompt_version.description or "",
        model_provider=prompt_version.model_provider,
        model_name=prompt_version.model_name,
        template=prompt_version.template,
        template_type=prompt_template_type,
        template_format=prompt_template_format,
        invocation_parameters=prompt_version.invocation_parameters,
        tools=prompt_version.tools,
        response_format=prompt_version.response_format,
    )


def _prompt_from_orm_prompt(orm_prompt: models.Prompt) -> Prompt:
    source_prompt_id = (
        str(GlobalID(PromptNodeType.__name__, str(orm_prompt.source_prompt_id)))
        if orm_prompt.source_prompt_id
        else None
    )
    return Prompt(
        id=str(GlobalID(PromptNodeType.__name__, str(orm_prompt.id))),
        source_prompt_id=source_prompt_id,
        name=orm_prompt.name,
        description=orm_prompt.description,
    )
