import os
from typing import Iterator
from dotenv import load_dotenv
from fastapi import APIRouter, Depends, Request

from langchain_huggingface import HuggingFaceEndpoint
from langchain_core.prompts import PromptTemplate

from libs.header_api_auth import get_api_key
from pydantic import BaseModel
from fastapi.responses import StreamingResponse
from langchain_ollama import ChatOllama, OllamaLLM

from libs.transformer.get_chat_gradio import get_chat_gradio
from libs.transformer.get_chat_transformer import get_chat_transformers

load_dotenv()
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN", )

os.environ["HUGGINGFACEHUB_API_TOKEN"] = HUGGINGFACEHUB_API_TOKEN

router = APIRouter(prefix="/get-chat-response", tags=["chat"])

class ChatInputForm(BaseModel):
    textInput: str
    repo_id: str
    prompt: str

@router.post("/")
async def get_chat_respone(body: ChatInputForm, api_key: str = Depends(get_api_key)):

    prompt = get_prompt(body.prompt)

    try:
        llm = OllamaLLM(
            model=body.repo_id,
            temperature=0.2,
            # huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
        )

        messages = [
            ("system", prompt),
            ("human", body.textInput)
        ]

        messages = [
            {"role": "system", "content": prompt},
            {"role": "user", "content": body.textInput},
        ]

        response = llm.stream(messages)

        # response = get_chat_gradio(body.textInput)

        # response = get_chat_transformers(messages)

        print(response)

        return StreamingResponse(get_response(response), media_type='text/event-stream')
    except Exception:
        return {"success": False, "status": Exception}

def get_response(response: Iterator[str]):
    for chunk in response:
        yield chunk


checkWritting = """You'll be provided with a text. Convert the text to standard English.
---------------
IMPORTANT:
- If the text is empty, do nothing.
- If the given text maintains grammatical accuracy, no suggestions are needed.
- If the given text is empty, do nothing.
- If the given text contains any errors in grammatical accuracy, provide the corrected text.
"""

template = """You are a helpful assistant. Do whatever user require. Response in markdown format."""
baiGiang = """Provide the given phrase in English. Provide the correct and popularly used English phrase along with its American IPA pronunciation and a brief explanation for it. Use the correct English phrase to create 4 example sentences along with the example IPA and brief meanings. Finally, suggest 4 similar English phrases with the correct English version, along with American IPA and their brief meanings.
Provie your response in markdown format"""

def get_prompt(prompt: str):
    prompts = {
        'template' : template,
        'checkWritting': checkWritting,
        'baiGiang': baiGiang
    }

    return prompts.get('template', template)