import asyncio
import time

import aiohttp
import requests

from app.conf.app_config import EmbeddingConfig, app_config


class EmbeddingClientManager:
    def __init__(self, config: EmbeddingConfig):
        self.config = config
        self.client = None

    def _get_url(self):
        return f"http://{self.config.host}:{self.config.port}"


    def init(self, wait_for_ready: bool = True):
        self.client = aiohttp.ClientSession()

    async def close(self):
        if self.client:
            await self.client.close()
            self.client = None

    async def aembed_documents(self, texts: list) -> list:
        if not self.client:
            self.init(wait_for_ready=False)

        url = f"{self._get_url()}/embed"
        payload = {
            "inputs": texts,
            "parameters": {"truncate": True}
        }

        async with self.client.post(url, json=payload) as response:
            result = await response.json()
            if isinstance(result, list):
                return result
            return result.get("embeddings", [])

    async def aembed_query(self, text: str) -> list:
        embeddings = await self.aembed_documents([text])
        return embeddings[0] if embeddings else []


embedding_client = EmbeddingClientManager(app_config.embedding)

if __name__ == "__main__":
    print("Testing EmbeddingClientManager...")
    
    async def test():
        try:
            embedding_client.init(wait_for_ready=True)
            print("Initialization successful")
            
            text = "What is Deep Learning?"
            print(f"Test text: {text}")
            
            result = await embedding_client.aembed_query(text)
            print(f"Embedding successful")
            print(f"Vector length: {len(result)}")
            print(f"First 10 values: {result[:10]}")
            
            await embedding_client.close()
        except Exception as e:
            print(f"Test failed: {str(e)}")
    
    asyncio.run(test())