import traceback
from typing import TypeVar, Generic

from fastapi import APIRouter, HTTPException
from fastapi.params import Header, Depends
from langchain_core.messages import BaseMessage
from langgraph.graph import MessagesState
from langgraph.graph.state import CompiledStateGraph
from openkosmos_core.auth import AuthToken
from openkosmos_core.setting import Setting
from starlette.responses import StreamingResponse

from openkosmos_ai.api.model import ChatCompletionResponse, FinishReason, \
    ModelInfoList, ModelInfo, ChatCompletionRequest, GraphProcessorConfig
from openkosmos_ai.common.model import RunConfig

INPUT_T = TypeVar("INPUT_T")
RESULT_T = TypeVar("RESULT_T")

log = Setting.logger("processor")


class GraphStreamProcessor(Generic[INPUT_T, RESULT_T]):

    def on_task_start(self, node_name: str, task_input: INPUT_T):
        return f"\n### {node_name} 开始 \n"

    def on_task_stop(self, node_name: str, task_result: RESULT_T):
        return f"\n### {node_name} 结束 \n"

    def on_message(self, node_name: str, meta_info: dict, message: BaseMessage):
        return message.content

    def on_create_agent_state(self, request: ChatCompletionRequest):
        return MessagesState(messages=[m.model_dump() for m in request.messages])


class AppInfo:

    def __init__(self, app: CompiledStateGraph,
                 stream_processor: GraphStreamProcessor = GraphStreamProcessor()):
        self.app = app
        self.stream_processor = stream_processor

    def name(self):
        return self.app.name

    def __str__(self):
        return f"{self.app.name}, stream_processor={self.stream_processor.__class__}"


class GraphProcessor:

    def __init__(self, config: GraphProcessorConfig = GraphProcessorConfig()):
        self._config = config
        self.app_map: dict[str, AppInfo] = {}
        self.model_info_list = ModelInfoList()
        if config.auth is not None:
            self._auth_token = AuthToken(config.auth)

    def config(self) -> GraphProcessorConfig:
        return self._config

    def add_app(self, app: AppInfo):
        log.info("add app: {}".format(app))
        self.app_map[app.name()] = app
        return self

    def agent_model_name(self):
        return self.model_info_list.model_dump()

    def get_model_list(self):
        return ModelInfoList(data=[ModelInfo(id=a.name()) for a in self.app_map.values()]).model_dump()

    def create_response(self, request: ChatCompletionRequest, config: RunConfig = None, context=None):
        current_app_info = self.app_map.get(request.model)
        if current_app_info is None:
            raise HTTPException(status_code=404,
                                detail=f"the model:'{request.model}' does not exist or you do not have access to it.")

        run_config = self.create_app_config(current_app_info, request) if config is None else config
        run_context = self.create_app_context(current_app_info, request) if context is None else context

        if request.stream:
            return StreamingResponse(
                self.stream_generator(current_app_info, request, run_config, run_context),
                media_type="text/event-stream")
        else:
            response = current_app_info.app.invoke(
                current_app_info.stream_processor.on_create_agent_state(request),
                context=run_context, config=run_config.to_config())

            return ChatCompletionResponse(model=current_app_info.name()).default_content(
                response["messages"][1].content).default_finish_reason(FinishReason.STOP).model_dump(
                exclude_none=self.config().response_config.exclude_none)

    def create_app_config(self, app_info: AppInfo, request: ChatCompletionRequest):
        return RunConfig(run_name=app_info.name())

    def create_app_context(self, app_info: AppInfo, request: ChatCompletionRequest):
        return None

    async def stream_generator(self, app_info: AppInfo, request: ChatCompletionRequest, config: RunConfig,
                               context):
        try:
            flow_result = app_info.app.astream(
                app_info.stream_processor.on_create_agent_state(request),
                config=config.to_config(),
                context=context,
                stream_mode=["messages", "tasks"])

            response_data = ChatCompletionResponse(model=app_info.name())

            yield response_data.get_message(response_config=self.config().response_config)

            async for mode, payload in flow_result:
                # print("\n---------------------------\n")
                # print(":", mode, payload)
                match mode:
                    case "tasks":
                        if payload.get("input") is None:
                            content = app_info.stream_processor.on_task_stop(
                                payload["name"], payload["result"])
                        else:
                            content = app_info.stream_processor.on_task_start(
                                payload["name"], payload["input"])
                        yield response_data.get_message(phase=payload["name"], content=content,
                                                        response_config=self.config().response_config)
                    case "messages":
                        message: BaseMessage = payload[0]
                        meta_info = payload[1]
                        # print("::", message, meta_info)
                        for content_blocks in message.content_blocks:
                            # print(":::", content_blocks)
                            match content_blocks["type"]:
                                case "text":
                                    yield response_data.get_message(
                                        content=app_info.stream_processor.on_message(
                                            meta_info["langgraph_node"], meta_info=meta_info,
                                            message=message),
                                        response_config=self.config().response_config)

            yield response_data.get_message(finish_reason=FinishReason.STOP, with_done=True,
                                            response_config=self.config().response_config)

        except Exception as e:
            traceback.print_exc()
            yield f"data: {str(e)} \n\n"

    async def _validate_auth_token(self, header_token: str = Header(default=None, alias="Authorization")):
        if self.config().auth is None:
            return None
        try:
            return self._auth_token.retrieve(header_token[6:])
        except:
            raise HTTPException(status_code=401, detail="invalid token")

    def create_api_router(self, tags: list[str] = None):
        _router = APIRouter(tags=self.config().tags if tags is None else tags)

        @_router.post("/chat/completions")
        async def _chat_completions(request: ChatCompletionRequest,
                                    token_payload=Depends(self._validate_auth_token)):
            return self.create_response(request)

        @_router.get("/models")
        async def _models(token_payload=Depends(self._validate_auth_token)):
            return self.get_model_list()

        @_router.post("/responses")
        async def _responses(request, token_payload=Depends(self._validate_auth_token)):
            pass

        return _router
