from typing import AsyncGenerator, Literal
from pydantic import BaseModel

from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.schema import BaseChatMessageHistory, Document
from langchain.schema.output_parser import StrOutputParser
from langchain.memory import ConversationBufferMemory
from langchain.schema.runnable import RunnableMap
from langchain.prompts import ChatPromptTemplate
from langchain.chains import OpenAIModerationChain


from config import Settings


class ChatRequest(BaseModel):
    session_id: str
    message: str


class ChatSSEResponse(BaseModel):
    type: Literal["context", "start", "streaming", "end", "error"]
    value: str | list[Document]


async def generate_standalone_question(
        chat_history: str,
        question: str,
        settings: Settings) -> str:
    """
    Generate question from the previous chats and the user prompt or question
    to get a stand alone question.

    :param chat_history: History of chats
    :param question: Question user wants to answer
    :param settings: Settings basemodel for openai keys
    :return: String of standalone question
    """

    prompt = PromptTemplate.from_template(
        template="""Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.
Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:""",
        input=[]
    )

    llm = ChatOpenAI(temperature=0, openai_api_key=settings.openai_api_key)

    chain = prompt | llm | StrOutputParser()  # type: ignore

    return await chain.ainvoke(  # type: ignore
        {
            "chat_history": chat_history,
            "question": question,
        }
    )


async def generate_sse_response(
        memory: ConversationBufferMemory,
        chat_memory: BaseChatMessageHistory,
        message: str,
        settings: Settings,
        prompt_template: ChatPromptTemplate,
        moderator_model: OpenAIModerationChain) -> AsyncGenerator[str, ChatSSEResponse]:
    """
    Function to generate the response using Server Sent Events(SSE)
    TODO: Generate good exception handling
    TODO: Adding the moderator to add a safeguard againt violations, hate speech and violent content

    :param context: List of all relevant documents with respect to user question to model
    :param chat_memory: Memory of previous conversations
    :param settings: Settings basemodel for openai keys
    :param prompt_template: The chat prompt template feed to the AI, with System messages and human messages
    :return: List of AsyncGenerator containing str and ChatSSEResponse
    """

    llm = ChatOpenAI(temperature=0, openai_api_key=settings.openai_api_key)

    # first get the memory then expand the memory
    # https://python.langchain.com/docs/guides/expression_language/cookbook
    chain = RunnableMap({
        "input": lambda x: x["input"],
        "memory": memory.load_memory_variables
    }) | {
        "input": lambda x: x["input"],
        "history": lambda x: x["memory"]["history"]
    } | prompt_template | llm

    response = ""

    try:
        yield ChatSSEResponse(type="start", value="").json()
        # type: ignore
        async for token in chain.astream({"input": message}):
            yield ChatSSEResponse(type="streaming", value=token.content).json()
            response += token.content

        yield ChatSSEResponse(type="end", value="").json()
        chat_memory.add_user_message(message=message)
        chat_memory.add_ai_message(message=response)

    except Exception as e:  # TODO: Add proper exception handling
        yield ChatSSEResponse(type="error", value=str(e)).json()

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from pydantic import BaseModel

from typing import AsyncGenerator


class ChatRequest(BaseModel):
    session_id: str
    message: str


app = FastAPI(
    title="QA Chatbot Streaming using FastAPI, LangChain Expression Language , OpenAI, and Chroma",
    version="0.1.0",
)


async def generate_response(
    context: str,
    message: str,
) -> AsyncGenerator[str, None]:
    prompt = PromptTemplate.from_template(
        """Answer the question based only on the following context:
{context}
Question: {question}"""
    )

    llm = ChatOpenAI(
        temperature=0,
        openai_api_key=YOUR_OPENAI_KEY,
    )

    chain = prompt | llm  # type: ignore

    response = ""
    async for token in chain.astream({"context": context, "question": message}):  # type: ignore
        yield token.content
        response += token.content


@app.post("/chat")
async def chat(
    request: ChatRequest,
) -> StreamingResponse:
    return StreamingResponse(
        generate_response(
            context="",
            message=request.message,
        ),
        media_type="text/plain",
    )