import json

from fastapi import FastAPI, HTTPException, Request
from pydantic import ValidationError
from sse_starlette import EventSourceResponse

from .activity_helpers import (
    async_sse_response_format,
    get_ignore_ping_comment,
    json_response_format,
)


try:
    import infiagent
    from infiagent.schemas import ChatCompleteRequest
    from infiagent.services.complete_local_test import (
        chat_local_event,
        chat_local_event_generator,
    )
    from infiagent.utils import get_logger
except ImportError:
    print("import infiagent failed, please install infiagent by 'pip install .' in the pipeline directory of ADA-Agent")
    from ..schemas import ChatCompleteRequest
    from ..services.complete_local_test import (
        chat_local_event,
        chat_local_event_generator,
    )
    from ..utils import get_logger

logger = get_logger()
local_app = FastAPI()


@local_app.post("/local_sse_test")
async def complete_sse(request: Request):
    body_str = await request.body()

    try:
        chat_request = ChatCompleteRequest.parse_raw(body_str)
        logger.info("Got chat request: {}".format(chat_request))
    except ValidationError as e:
        error_msg = "Invalid input chat_request. Error: {}".format(str(e))
        raise HTTPException(status_code=500, detail=error_msg)

    return EventSourceResponse(async_sse_response_format(chat_local_event_generator(chat_request)),
                               ping_message_factory=get_ignore_ping_comment())


@local_app.post("/local_json_test")
async def complete_json(request: Request):

    body_str = await request.body()

    try:
        chat_request = ChatCompleteRequest.parse_raw(body_str)
        logger.info("Got chat request: {}".format(chat_request))
    except ValidationError as e:
        error_msg = "Invalid input chat_request. Error: {}".format(str(e))
        raise HTTPException(status_code=500, detail=error_msg)

    response_items = await chat_local_event(chat_request)
    return json_response_format(response_items)


@local_app.post("/exception_test")
async def complete_json(request: Request):
    body_str = await request.body()

    try:
        chat_request = ChatCompleteRequest.parse_raw(body_str)
        logger.info("Got chat request: {}".format(chat_request))
    except ValidationError as e:
        error_msg = "Invalid input chat_request. Error: {}".format(str(e))
        raise HTTPException(status_code=500, detail=error_msg)
    return EventSourceResponse(async_sse_response_format(chat_local_event_generator(chat_request)))


async def exception_test(request: Request):
    body_str = await request.body()
    json_val = json.loads(body_str)
    exception_type = json_val.get("exception", None)

    if exception_type:
        raise ValueError("Error triggerd!")
    else:
        yield iter(["Success"])