#!/usr/bin/python3
# -*- coding: utf-8 -*-
from __future__ import annotations

import asyncio
import contextlib
import pathlib
import shutil
import traceback
import uuid
from collections import deque
from functools import partial
from json import JSONDecodeError
from typing import Dict

import fire
import openai
import tenacity
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from loguru import logger
from metagpt.config import CONFIG
from metagpt.logs import set_llm_stream_logfunc
from metagpt.schema import Message
from openai import OpenAI

from data_model import (
    NewMsg,
    MessageJsonModel,
    Sentences,
    Sentence,
    SentenceType,
    SentenceValue,
    ThinkActPrompt,
    LLMAPIkeyTest,
    ThinkActStep,
)
from message_enum import QueryAnswerType, MessageStatus
from software_company import RoleRun, SoftwareCompany


class Service:
    @classmethod
    async def create_message(cls, req_model: NewMsg, request: Request):
        """
        Session message stream
        """
        tc_id = 0
        task = None
        try:
            exclude_keys = CONFIG.get("SERVER_METAGPT_CONFIG_EXCLUDE", [])
            config = {k.upper(): v for k, v in req_model.config.items() if k not in exclude_keys}
            cls._set_context(config)
            msg_queue = deque()
            CONFIG.LLM_STREAM_LOG = lambda x: msg_queue.appendleft(x) if x else None

            role = SoftwareCompany()
            role.recv(message=Message(content=req_model.query))
            answer = MessageJsonModel(
                steps=[
                    Sentences(
                        contents=[
                            Sentence(
                                type=SentenceType.TEXT.value, value=SentenceValue(answer=req_model.query), is_finished=True
                            )
                        ],
                        status=MessageStatus.COMPLETE.value,
                    )
                ],
                qa_type=QueryAnswerType.Answer.value,
            )

            async def stop_if_disconnect():
                while not await request.is_disconnected():
                    await asyncio.sleep(1)

                if task is None:
                    return

                if not task.done():
                    task.cancel()
                    logger.info(f"cancel task {task}")

            asyncio.create_task(stop_if_disconnect())

            while True:
                tc_id += 1
                if await request.is_disconnected():
                    return
                think_result: RoleRun = await role.think()
                if not think_result:  # End of conversion
                    break

                think_act_prompt = ThinkActPrompt(role=think_result.role.profile)
                think_act_prompt.update_think(tc_id, think_result)
                yield think_act_prompt.prompt + "\n\n"
                task = asyncio.create_task(role.act())

                while not await request.is_disconnected():
                    if msg_queue:
                        think_act_prompt.update_act(msg_queue.pop(), False)
                        yield think_act_prompt.prompt + "\n\n"
                        continue

                    if task.done():
                        break

                    await asyncio.sleep(0.5)
                else:
                    task.cancel()
                    return

                act_result = await task
                think_act_prompt.update_act(act_result)
                yield think_act_prompt.prompt + "\n\n"
                answer.add_think_act(think_act_prompt)
            yield answer.prompt + "\n\n"  # Notify the front-end that the message is complete.
        except asyncio.CancelledError:
            task.cancel()
        except tenacity.RetryError as retry_error:
            yield cls.handle_retry_error(tc_id, retry_error)
        except Exception as ex:
            description = str(ex)
            answer = traceback.format_exc()
            think_act_prompt = cls.create_error_think_act_prompt(tc_id, description, description, answer)
            yield think_act_prompt.prompt + "\n\n"
        finally:
            CONFIG.WORKSPACE_PATH: pathlib.Path
            if CONFIG.WORKSPACE_PATH.exists():
                shutil.rmtree(CONFIG.WORKSPACE_PATH)

    @staticmethod
    def create_error_think_act_prompt(tc_id: int, title, description: str, answer: str) -> ThinkActPrompt:
        step = ThinkActStep(
            id=tc_id,
            status="failed",
            title=title,
            description=description,
            content=Sentence(type=SentenceType.ERROR.value, id=1, value=SentenceValue(answer=answer), is_finished=True),
        )
        return ThinkActPrompt(step=step)

    @classmethod
    def handle_retry_error(cls, tc_id: int, error: tenacity.RetryError):
        # Known exception handling logic
        try:
            # Try to get the original exception
            original_exception = error.last_attempt.exception()
            while isinstance(original_exception, tenacity.RetryError):
                original_exception = original_exception.last_attempt.exception()

            if isinstance(original_exception, openai.AuthenticationError):
                answer = original_exception.message
                title = "OpenAI AuthenticationError"
                think_act_prompt = cls.create_error_think_act_prompt(tc_id, title, title, answer)
                return think_act_prompt.prompt + "\n\n"
            elif isinstance(original_exception, openai.APITimeoutError):
                answer = original_exception.message
                title = "OpenAI APITimeoutError"
                think_act_prompt = cls.create_error_think_act_prompt(tc_id, title, title, answer)
                return think_act_prompt.prompt + "\n\n"
            elif isinstance(original_exception, JSONDecodeError):
                answer = str(original_exception)
                title = "MetaGPT Error"
                description = "LLM return result parsing error"
                think_act_prompt = cls.create_error_think_act_prompt(tc_id, title, description, answer)
                return think_act_prompt.prompt + "\n\n"
            else:
                return cls.handle_unexpected_error(tc_id, error)
        except Exception:
            return cls.handle_unexpected_error(tc_id, error)

    @classmethod
    def handle_unexpected_error(cls, tc_id, error):
        description = str(error)
        answer = traceback.format_exc()
        think_act_prompt = cls.create_error_think_act_prompt(tc_id, description, description, answer)
        return think_act_prompt.prompt + "\n\n"

    @staticmethod
    def _set_context(context: Dict) -> Dict:
        uid = uuid.uuid4().hex
        context["WORKSPACE_PATH"] = pathlib.Path("workspace", uid)
        for old, new in (("DEPLOYMENT_ID", "DEPLOYMENT_NAME"), ("OPENAI_API_BASE", "OPENAI_BASE_URL")):
            if old in context and new not in context:
                context[new] = context[old]
        CONFIG.set_context(context)
        return context


default_llm_stream_log = partial(print, end="")


def llm_stream_log(msg):
    with contextlib.suppress():
        CONFIG._get("LLM_STREAM_LOG", default_llm_stream_log)(msg)


class ChatHandler:
    @staticmethod
    async def create_message(req_model: NewMsg, request: Request):
        """Message stream, using SSE."""
        event = Service.create_message(req_model, request)
        headers = {"Cache-Control": "no-cache", "Connection": "keep-alive"}
        return StreamingResponse(event, headers=headers, media_type="text/event-stream")


class LLMAPIHandler:
    @staticmethod
    async def check_openai_key(req_model: LLMAPIkeyTest):
        try:
            # Listing all available models.
            client = OpenAI(api_key=req_model.api_key)
            response = client.models.list()
            model_set = {model.id for model in response.data}
            if req_model.llm_type in model_set:
                logger.info("API Key is valid.")
                return JSONResponse({"valid": True})
            else:
                logger.info("API Key is invalid.")
                return JSONResponse({"valid": False, "message": "Model not found"})
        except Exception as e:
            # If the request fails, return False
            logger.info(f"Error: {e}")
            return JSONResponse({"valid": False, "message": str(e)})


app = FastAPI()

app.mount(
    "/storage",
    StaticFiles(directory="./storage/"),
    name="storage",
)

app.add_api_route(
    "/api/messages",
    endpoint=ChatHandler.create_message,
    methods=["post"],
    summary="Session message sending (streaming response)",
)
app.add_api_route(
    "/api/test-api-key",
    endpoint=LLMAPIHandler.check_openai_key,
    methods=["post"],
    summary="LLM APIkey detection",
)

app.mount(
    "/",
    StaticFiles(directory="./static/", html=True, follow_symlink=True),
    name="static",
)

set_llm_stream_logfunc(llm_stream_log)


def main():
    server_config = CONFIG.get("SERVER_UVICORN", {})
    uvicorn.run(app="__main__:app", **server_config)


if __name__ == "__main__":
    fire.Fire(main)