import aiohttp
import asyncio
from typing import List, Callable, Optional

class HTTPEmbeddingClient():
    def __init__(
        self,
        host: str,
        port: Optional[int] = None,
        endpoint: str = "/",
        payload_builder: Optional[Callable[[str], dict]] = None,
        trust_env: bool = True,
    ) -> None:
        self.host = host
        self.port = port
        self.endpoint = endpoint.rstrip("/") or "/"
        self.payload_builder = payload_builder or (lambda text: {"text": text})
        self.trust_env = trust_env

        if self.port is not None:
            self.base_url = f"http://{self.host}:{self.port}"
        else:
            self.base_url = f"http://{self.host}"

    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(trust_env=self.trust_env) as session:
            async with session.post(f"{self.base_url}{self.endpoint}", json=self.payload_builder(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(self.endpoint, json=self.payload_builder(text))
            return await response.json(content_type=None)

        tasks = []
        async with aiohttp.ClientSession(self.base_url, trust_env=self.trust_env) 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)

# Backward-compatible alias
EmbeddingClient = HTTPEmbeddingClient


if __name__=="__main__":
    import asyncio, time

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

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

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

    asyncio.run(main())