import json, requests


class YandexGPTEntryPoint:
    def __init__(self, token, folder_id):
        iam_token = requests.post(
            url='https://iam.api.cloud.yandex.net/iam/v1/tokens', 
            data=json.dumps({'yandexPassportOauthToken': token})
            ).json()['iamToken']
        self.folder_id = folder_id
        self.text_url = 'https://llm.api.cloud.yandex.net/foundationModels/v1/completion'
        self.emb_url = 'https://llm.api.cloud.yandex.net:443/foundationModels/v1/textEmbedding'
        self.doc_uri = f'emb://{self.folder_id}/text-search-doc/latest'
        self.query_uri = f'emb://{self.folder_id}/text-search-query/latest'
        self.headers = {
            'Authorization': f'Bearer {iam_token}', 
            'Content-Type': 'application/json'
            }
        self.DIM = 256
        self.ZEROS = [0 for _ in range(self.DIM)]
        self.ERROR_MESSAGE = ''
        self.warmup()

    def get_response(self, total_input: str) -> str:
        data = {
            'model_uri': f'gpt://{self.folder_id}/yandexgpt/latest',
            'messages': [
                {
                    'role': 'user',
                    'text': total_input,
                }
            ],
        }
        try:
            response = requests.post(
                url=self.text_url, 
                headers=self.headers, 
                json=data
                ).json()
            return response['result']['alternatives'][0]['message']['text']
        except:
            return self.ERROR_MESSAGE

    def get_embeddings(self, total_input, input_is_long=False) -> list:
        data = {
            'modelUri': self.doc_uri if input_is_long else self.query_uri, 
            'text': total_input
            }
        try:
            response = requests.post(
                url=self.emb_url, 
                headers=self.headers, 
                json=data
                ).json()
            return response['embedding']
        except:
            return self.ZEROS

    def warmup(self) -> None:
        assert self.get_response('Прогрев') != self.ERROR_MESSAGE, 'Нет доступа к ллм!'
        assert self.get_embeddings('Прогрев') != self.ZEROS, 'Нет доступа к ллм!'
        assert self.get_embeddings('Прогрев', input_is_long=True) != self.ZEROS, 'Нет доступа к ллм!'
                