"""
@author: Qichang Zheng
@email: qichangzheng@uchicago.edu
@date: 2023-11-11
@personal website: http://qichangzheng.net
This is a personal API developed by Qichang Zheng. If you have any questions, please feel free to contact me via email.
"""

import requests
from time import sleep
from ping3 import ping
from bs4 import BeautifulSoup
import os
from tqdm import tqdm

__version__ = '0.0.14'

app_api_dict = {
    "GPT3.5": 'fastgpt-ZadxgXx4pZV4OqSyi7dcYHambSq',
    'GPT4': 'fastgpt-BmDbL0SAmNUcyrRxDK5SF3lApP',
    'stock_api': 'fastgpt-lmiAMrgderiIfidhxQ5HVOD6PJIr3ntpxO4',
    'shannon_test1': 'fastgpt-I8yveuOpvtf5NNjCc4feXkscEQGiKxwH3J'
}

server_dict = {
    'Virginia': '54.159.212.206',
    'Singapore': '47.236.36.36'
}

fastgpt_server = {
    'Virginia': 'http://fastgpt.qichangzheng.net/api/v1/chat/completions',
    'Singapore': 'http://nbgpt.qichangzheng.net/api/v1/chat/completions'
}

class LLM_API:
    def __init__(self):
        self.app_api_dict = {
            "GPT3.5": 'fastgpt-ZadxgXx4pZV4OqSyi7dcYHambSq',
            'GPT4': 'fastgpt-BmDbL0SAmNUcyrRxDK5SF3lApP',
            'stock_api': 'fastgpt-lmiAMrgderiIfidhxQ5HVOD6PJIr3ntpxO4',
            'shannon_test1': 'fastgpt-I8yveuOpvtf5NNjCc4feXkscEQGiKxwH3J',
        }
        self.server_dict = {
            'Virginia': '54.159.212.206',
            'Singapore': '47.236.36.36',
        }
        try:
            self.server = self.select_server()
        except:
            print(f'Auto server selection failed, using default server Virginia, '
                  f'you can also select server manually by setting self.server == "Singapore",'
                  f'available servers: {list(self.server_dict.keys())}')
            self.server = 'Virginia'


    def select_server(self):
        # Initialize a dictionary to store ping results
        ping_results = {}

        # Ping each server and store the results
        for server_name, server_ip in server_dict.items():
            ping_result = ping(server_ip)
            if ping_result is not None:
                ping_results[server_name] = ping_result

        # Find the server with the lowest ping
        lowest_ping_server = min(ping_results, key=ping_results.get)
        lowest_ping = ping_results[lowest_ping_server]

        # Print the ping results
        print("Ping")
        for server_name, ping_time in ping_results.items():
            print(f"{server_name}: {ping_time:.2f} ms")

        # Print the server with the lowest ping
        print(f"{lowest_ping_server} with lowest ping ({lowest_ping:.2f} ms) selected.")
        return lowest_ping_server


    def chat(self, app_or_key, message, chatId=None):
        if app_or_key.startswith('fastgpt-'):
            apikey = app_or_key
        else:
            try:
                apikey = app_api_dict[app_or_key]
            except KeyError:
                raise KeyError(f'App {app_or_key} not found, available apps are {list(app_api_dict.keys())}')
        url = fastgpt_server[self.server]
        headers = {
            "Authorization": 'Bearer ' + apikey,
            "Content-Type": "application/json"
        }
        data = {
            "chatId": chatId,
            "stream": False,
            "detail": False,
            "variables": {
                "cTime": "2023/10/18 22:22"
            },
            "messages": [
                {
                    "content": message,
                    "role": "user"
                }
            ]
        }
        while True:
            try:
                response = requests.post(url, headers=headers, json=data).json()['choices'][0]['message']['content']
                break
            except:
                sleep(3)
        return response


class Model_Downloader:
    def __init__(self):
        self.revision = 'revision_not_found'
        print("You may use the downloader by: Model_Downloader.download(model_name, path)")


    def get_redirect_urls(self, urls):
        result = []
        for url in urls:
            res = requests.get(url, allow_redirects=False)
            result.append('https://cdn.qichangzheng.net' + res.text.split('huggingface.co')[1])
        return result

    def get_file_urls(self, model):
        # Base URL for the model directory
        base_url = f"https://huggingface.qichangzheng.net/{model}/tree/main"

        # Send a GET request to fetch the page content
        response = requests.get(base_url)
        response.raise_for_status()  # Raises an HTTPError if the HTTP request returned an unsuccessful status code

        # Parse the HTML content of the page
        soup = BeautifulSoup(response.content, 'html.parser')
        self.revision = soup.find('a', class_='rounded').text

        # Find all the <a> tags with the 'download' attribute
        direct_urls = []
        redirect_urls = []
        for tag in soup.find_all('a', attrs={'title': 'Download file'}):
            if tag.find('svg').next_sibling.strip() == '':
                direct_urls.append('https://huggingface.qichangzheng.net' + tag['href'])
            else:
                redirect_urls.append('https://huggingface.qichangzheng.net' + tag['href'])
        direct_urls.extend(self.get_redirect_urls(redirect_urls))
        return direct_urls

    def download(self, model_name, path):
        # Ensure the download directory exists
        urls = self.get_file_urls(model_name)
        path = path + '/models--' + '--'.join(model_name.split('/')) + f'/{self.revision}'
        print(f'Model will be saved at {path}')
        if not os.path.exists(path):
            os.makedirs(path)

        for url in tqdm(urls, desc="Downloading files", unit="file"):
            # Extract the file name from the URL
            filename = url.split('/')[-1].split('?')[0]  # Assumes the URL ends with filename?download=true
            full_path = os.path.join(path, filename)

            # Stream the download to avoid using too much memory
            response = requests.get(url, stream=True)

            # Check if the request was successful
            if response.status_code == 200:
                # Open the file to write as binary - 'wb'
                with open(full_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=1024):
                        if chunk:  # filter out keep-alive new chunks
                            f.write(chunk)
            else:
                print(f"Failed to download {url}")


