import aiohttp
import asyncio
import requests
from typing import List

class EmbeddingClient():
    def __init__(self, host: str, port: int) -> None:
        self.host = host
        self.port = port
        self.loop = asyncio.get_event_loop()

    async def get_embedding(self, text: str) -> 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.
        
        Returns:
            dict: The embedding of the given text in the form of a dictionary.
        """
        async with aiohttp.ClientSession() as session:
            async with session.post(f"http://{self.host}:{self.port}/", json={'text':text}) as response:
                return await response.json()

    async def _abatch_get_embedding(self, texts) -> List[dict]:

        async def _a_get_embedding(session, text) -> dict:
            response = await session.post('/', json={"text": text})
            return await response.json(content_type=None)
        
        tasks = []
        async with aiohttp.ClientSession(f"http://{self.host}:{self.port}/") as session:
            for text in texts:
                tasks.append(_a_get_embedding(session, text))

            result = await asyncio.gather(*tasks)
        return result
    
    async def batch_get_embedding(self, texts):
         return await self._abatch_get_embedding(texts)

    
if __name__=="__main__":
    import time

    emb_client = EmbeddingClient(host="localhost", port=12081)
    text = "พนักงานทุกท่านซึ่งปฏิบัติงานมาครบ 1ปีจะได้รับวันหยุดประจำปี ดังนี้\n\nวันหยุดพักผ่อนประจำปี\n\nพนักงานทุกท่านจะได้รับวันหยุดพักผ่อนประจำปีท่านละ  10 วัน สำหรับพนักงานใหม่ที่ผ่านการทดลองงานแล้วก็ให้ลดลงตามส่วน\n\nถ้าลาเกินสิทธิ 10 วัน"

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

    start = time.time()
    for i in range(5):
        res = emb_client.get_embedding(text)
    # print(res)
    print("sync time:", time.time()-start)
    print("--"*20)
    start = time.time()
    result = emb_client.batch_get_embedding([text for i in range(5)])
    print(len(result))
    print("async time:", time.time()-start)