import re
import time
import logging
import aiohttp
import requests
import asyncio

from .settings import TIDB_AI_API_BASE

logger = logging.getLogger(__name__)


def retry_on_exception(retry_times=3, exception=Exception, sleep_time=0):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for i in range(retry_times):
                try:
                    return func(*args, **kwargs)
                except exception as e:
                    time.sleep(sleep_time)
                    continue
        return wrapper
    return decorator


def remove_citation(content:str) -> str:
    # remove citation like [^1] or [^1^]
    return re.sub(r'\[\^(\d+)\^*\]', '', content)


class CommunityQAService:
    def __init__(self, api_key: str, chat_engine: str):
        self.api_base = TIDB_AI_API_BASE
        self.chat_endpoint = self.api_base + "/chats"
        self.feedback_endpoint = self.api_base + "/chat-messages/{chat_message_id}/feedback"
        self.api_key = api_key
        self.chat_engine = chat_engine

    @retry_on_exception(retry_times=3, exception=Exception, sleep_time=1)
    def ask(self, origin: str, qa_context:list) -> tuple[str, str, str]:
        response = requests.post(
            self.chat_endpoint,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}",
                "Origin": origin,
            },
            json={
                "messages": qa_context,
                "index": "default",
                "chat_engine": self.chat_engine,
                "stream": False,
            },
            timeout=10000,
        )
        response.raise_for_status()
        data = response.json()
        return remove_citation(data["content"]), data["chat_id"], data["message_id"]

    @retry_on_exception(retry_times=3, exception=Exception, sleep_time=1)
    async def ask_async(self, origin: str, qa_context:list) -> tuple[str, str, str]:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                self.chat_endpoint,
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {self.api_key}",
                    "Origin": origin,
                },
                json={
                    "messages": qa_context,
                    "index": "default",
                    "chat_engine": self.chat_engine,
                    "stream": False,
                },
                timeout=10000,
            ) as response:
                response.raise_for_status()
                data = await response.json()
                return remove_citation(data["content"]), data["chat_id"], data["message_id"]

    async def _feedback(self, origin: str, chat_message_id: str, username: str, feedback_type: str, comment: str):
        async with aiohttp.ClientSession() as session:
            async with session.post(
                self.feedback_endpoint.format(chat_message_id=chat_message_id),
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {self.api_key}",
                    "Origin": origin,
                },
                json={
                    "feedback_type": feedback_type,
                    "comment": comment,
                },
                timeout=1000,
            ) as response:
                response.raise_for_status()

    async def like_async(self, origin: str, chat_message_id: str, username: str):
        return await self._feedback(origin, chat_message_id, username, "like", f"From {origin}, user {username}")

    async def dislike_async(self, origin: str, chat_message_id: str, username: str, comment: str):
        return await self._feedback(origin, chat_message_id, username, "dislike", comment)

    def like(self, origin: str, chat_message_id: str, username: str):
        return asyncio.run(self.like_async(origin, chat_message_id, username))

    def dislike(self, origin: str, chat_message_id: str, username: str, comment: str):
        return asyncio.run(self.dislike_async(origin, chat_message_id, username, comment))


class CommunityNewQaService:
    def __init__(self):
        self.chat_endpoint = "https://pingcap-doc-search.htapdb.com/ask"

    @retry_on_exception(retry_times=3, exception=Exception, sleep_time=1)
    def ask(self, origin: str, qa_context:list) -> tuple[str, str, str]:
        response = requests.post(
            self.chat_endpoint,
            headers={
                "Content-Type": "application/json",
                "Origin": origin,
            },
            json={
                "messages": qa_context,
            },
            timeout=10000,
        )
        response.raise_for_status()
        data = response.json()
        return remove_citation(data["content"]), data["chat_id"], data["message_id"]

    @retry_on_exception(retry_times=3, exception=Exception, sleep_time=1)
    async def ask_async(self, origin: str, qa_context:list) -> tuple[str, str, str]:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                self.chat_endpoint,
                headers={
                    "Content-Type": "application/json",
                    "Origin": origin,
                },
                json={
                    "messages": qa_context,
                },
                timeout=10000,
            ) as response:
                response.raise_for_status()
                data = await response.json()
                return remove_citation(data["content"]), data["chat_id"], data["message_id"]