import os
import requests
import base64
from kbrainsdk.admin import Admin
from kbrainsdk.llms import LLMs
from kbrainsdk.datasets import Datasets
from kbrainsdk.ingest import Ingest
from kbrainsdk.ai import AI
from kbrainsdk.users import User
from kbrainsdk.events import Events
from kbrainsdk.signalr import SignalR

VALID_MODEL_NAMES = ["gpt-3.5-turbo", "gpt-4", "gpt-3.5-instruct"]

class KBRaiNAPI:
    def __init__(self, base_url=None, account_id=None, api_key=None):
        self.KBRAIN_NAME = "KBRaiN ™"
        self.base_url = base_url or os.getenv('KBRAIN_BASE_URL')
        account_id = account_id or os.getenv('KBRAIN_ACCOUNT_ID')
        api_key = api_key or os.getenv('KBRAIN_API_KEY')
        auth_string = f"{account_id}:{api_key}"
        auth_bytes = base64.b64encode(auth_string.encode('utf-8')).decode('utf-8')
        self.headers = {'Authorization': f'Basic {auth_bytes}'}
        self.admin = Admin(self)
        self.llms = LLMs(self)
        self.datasets = Datasets(self)
        self.ingest = Ingest(self)
        self.ai = AI(self)
        self.users = User(self)
        self.events = Events(self)
        self.signalr = SignalR(self)
        
    def healthy(self):
        path = f"/health/v1"
        response = self.call_endpoint(path, {}, "get")
        return response

    
    def call_endpoint(self, path, payload, method):
        request_method = getattr(requests, method)
        endpoint = f"{self.base_url}{path}"
        response = request_method(url=endpoint, json=payload, headers=self.headers)
        if response.status_code >= 200 and response.status_code <= 300:
            return response.json()
        
        message = response.json().get('message')
        raise Exception(message)