import aiohttp
import json
import time
from typing import List

class ModalEmbeddingClient():
    def __init__(self) -> None:
        self.web_url = "https://knight-h--bge-m3-bm25-model-get-embedding.modal.run"

    async def _get_embedding(self, text: str, is_query: bool = False) -> dict:
        """
        Sends a POST request to the embedding server to get the embedding for the given text.

        Args:
            text (str): The text to get the embedding for.
            is_query (bool): Whether the text is a query or not.

        Returns:
            dict: The embedding of the given text in the form of a dictionary.
        """
        start = time.time()
        async with aiohttp.ClientSession() as session:
            async with session.post(
                self.web_url,
                json={
                    "input": json.dumps({"doc_id":1, "document_text": text, "is_query": is_query}) + "\n",
                },
            ) as response:
                response_json = await response.json()

        print(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Modal serverless embedding latency",
            "latency": str(time.time() - start),
            "exception": False
        }))
        data = json.loads(response_json)
        # Get the embedding from the response
        data["embedding"] = data["dense_embeddings"] # keep forgetting to rename the keysss
        data["sparse_embedding"] = data["sparse_embeddings"]
        del data["dense_embeddings"], data["sparse_embeddings"]
        return data

    async def get_embedding(self, text: str, is_query: bool = False) -> dict:
        try:
            result = await self._get_embedding(text, is_query=is_query)
        except Exception as e:
            print(json.dumps({
                "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
                "message": "ERROR Modal serverless embedding",
                "latency": None,
                "exception": True
            }))
            raise(e)
        return result

    async def batch_get_embedding(self, texts: List[str]) -> List[dict]:
        """
        Sends a POST request to the embedding server to get embeddings for a list of texts.
        Returns a list of embedding dicts.
        """
        start = time.time()
        lines = [
            json.dumps({"doc_id": i, "document_text": text}, ensure_ascii=False)
            for i, text in enumerate(texts)
        ]
        jsonl_body = "\n".join(lines) + "\n"
        async with aiohttp.ClientSession() as session:
            async with session.post(
                self.web_url,
                json={
                    "input": jsonl_body
                },
            ) as response:
                response_json = await response.json()

        print(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Modal serverless batch embedding latency",
            "latency": str(time.time() - start),
            "exception": False
        }))
        lines = response_json.strip().split('\n')
        data_list = []
        for line in lines:
            data = json.loads(line)
            data["embedding"] = data["dense_embeddings"]
            data["sparse_embedding"] = data["sparse_embeddings"]
            del data["dense_embeddings"], data["sparse_embeddings"]
            data_list.append(data)
        return data_list

if __name__=="__main__":
    import asyncio
    emb_client = ModalEmbeddingClient()
    text = "พนักงานทุกท่านซึ่งปฏิบัติงานมาครบ 1ปีจะได้รับวันหยุดประจำปี ดังนี้\n\nวันหยุดพักผ่อนประจำปี\n\nพนักงานทุกท่านจะได้รับวันหยุดพักผ่อนประจำปีท่านละ  10 วัน สำหรับพนักงานใหม่ที่ผ่านการทดลองงานแล้วก็ให้ลดลงตามส่วน\n\nถ้าลาเกินสิทธิ 10 วัน"

    res = asyncio.run(emb_client.get_embedding(text))
    assert "embedding" in res
    print(res)
