"""
Official API for ChatGPT
"""
import asyncio
import json
import os
import sys

import httpx
import requests
import tiktoken
from OpenAIAuth.OpenAIAuth import OpenAIAuth

ENCODER = tiktoken.get_encoding("gpt2")


def get_max_tokens(prompt: str) -> int:
    """
    Get the max tokens for a prompt
    """
    return 4000 - len(ENCODER.encode(prompt))


class Message:
    """
    A single exchange between the user and the bot
    """

    def __init__(self, text: str, author: str) -> None:
        self.text: str = text
        self.author: str = author


class Conversation:
    """
    A single conversation
    """

    def __init__(self) -> None:
        self.messages: list[Message] = []


CONVERSATION_BUFFER: int = int(os.environ.get("CONVERSATION_BUFFER") or 1500)


class Conversations:
    """
    Conversation handler
    """

    def __init__(self) -> None:
        self.conversations: dict[str][Conversation] = {}

    def add_message(self, message: Message, conversation_id: str) -> None:
        """
        Adds a message to a conversation
        """
        if conversation_id not in self.conversations:
            self.conversations[conversation_id] = Conversation()
        self.conversations[conversation_id].messages.append(message)

    def get(self, conversation_id: str) -> str:
        """
        Builds a conversation string from a conversation id
        """
        if conversation_id not in self.conversations:
            return ""
        # Build conversation string from messages and check if it's too long
        conversation = ""
        for message in self.conversations[conversation_id].messages:
            conversation += f"{message.author}: {message.text}<|im_sep|>\n\n"
        if len(ENCODER.encode(conversation)) > 4000 - CONVERSATION_BUFFER:
            self.purge_history(conversation_id)
            return self.get(conversation_id)
        return conversation

    def purge_history(self, conversation_id: str, num: int = 1):
        """
        Remove oldest messages from a conversation
        """
        if conversation_id not in self.conversations:
            return
        self.conversations[conversation_id].messages = self.conversations[
            conversation_id
        ].messages[num:]

    def rollback(self, conversation_id: str, num: int = 1):
        """
        Remove latest messages from a conversation
        """
        if conversation_id not in self.conversations:
            return
        self.conversations[conversation_id].messages = self.conversations[
            conversation_id
        ].messages[:-num]

    def remove(self, conversation_id: str) -> None:
        """
        Removes a conversation
        """
        if conversation_id in self.conversations:
            del self.conversations[conversation_id]


BASE_PROMPT = (
    os.environ.get("BASE_PROMPT")
    or """You are ChatGPT, a large language model by OpenAI. Respond conversationally\n\n\n"""
)

PROXY_URL = os.environ.get("PROXY_URL") or "https://chat.duti.tech"


class Chatbot:
    """
    Handles everything seamlessly
    """

    def __init__(
        self,
        email: str,
        password: str,
        paid: bool = False,
        proxy=None,
        insecure: bool = False,
        session_token: str = None,
    ) -> None:
        self.proxy = proxy
        self.email: str = email
        self.password: str = password
        self.session_token = session_token
        self.insecure: bool = insecure
        self.api_key: str
        self.paid: bool = paid
        self.conversations = Conversations()
        self.login(email, password, proxy, insecure, session_token)

    async def ask(self, prompt: str, conversation_id: str = None) -> dict:
        """
        Gets a response from the API
        """
        if conversation_id is None:
            conversation_id = "default"
        self.conversations.add_message(
            Message(prompt, "User"),
            conversation_id=conversation_id,
        )
        conversation: str = self.conversations.get(conversation_id)
        # Build request body
        body = self.__get_config()
        body["prompt"] = BASE_PROMPT + conversation + "ChatGPT: "
        body["max_tokens"] = get_max_tokens(conversation)
        async with httpx.AsyncClient(proxies=self.proxy if self.proxy else None).stream(
            method="POST",
            url=PROXY_URL + "/completions",
            data=json.dumps(body),
            headers={"Authorization": f"Bearer {self.api_key}"},
            timeout=1080,
        ) as response:
            full_result = ""
            async for line in response.aiter_lines():
                if response.status_code == 429:
                    print("error: " + "Too many requests")
                    raise Exception("Too many requests")
                elif response.status_code == 523:
                    print("error: " + "Origin is unreachable. Ensure that you are authenticated and are using the correct pricing model.")
                    raise Exception("Origin is unreachable. Ensure that you are authenticated and are using the correct pricing model.")
                elif response.status_code == 503:
                    print("error: " + "OpenAI error!")
                    raise Exception("OpenAI error!")
                elif response.status_code != 200:
                    print("error: " + "Unknown error")
                    raise Exception("Unknown error")
                line = line.strip()
                if line == "\n" or line == "":
                    continue
                if line == "data: [DONE]":
                    break
                try:
                    # Remove "data: " from the start of the line
                    data = json.loads(line[6:])
                    if data is None:
                        continue
                    full_result += data["choices"][0]["text"].replace("<|im_end|>", "")
                    if "choices" not in data:
                        continue
                    yield data
                except json.JSONDecodeError:
                    continue
            self.conversations.add_message(
                Message(full_result, "ChatGPT"),
                conversation_id=conversation_id,
            )

    def __get_config(self) -> dict:
        return {
            "temperature": float(os.environ.get("TEMPERATURE") or 0.5),
            "top_p": float(os.environ.get("TOP_P") or 1),
            "stop": ["<|im_end|>", "<|im_sep|>"],
            "presence_penalty": float(os.environ.get("PRESENCE_PENALTY") or 1.0),
            "paid": self.paid,
            "stream": True,
        }

    def login(self, email, password, proxy, insecure, session_token) -> None:
        """
        Login to the API
        """
        if not insecure:
            auth = OpenAIAuth(email_address=email, password=password, proxy=proxy)
            if session_token:
                auth.session_token = session_token
                auth.get_access_token()
                self.api_key = auth.access_token
                if self.api_key is None:
                    self.session_token = None
                    self.login(email, password, proxy, insecure, None)
                return
            auth.begin()
            self.session_token = auth.session_token
            self.api_key = auth.access_token
        else:
            auth_request = requests.post(
                PROXY_URL + "/auth",
                json={"email": email, "password": password},
                timeout=10,
            )
            self.api_key = auth_request.json()["accessToken"]


def get_input(prompt):
    """
    Multi-line input
    """
    # Display the prompt
    print(prompt, end="")

    # Initialize an empty list to store the input lines
    lines = []

    # Read lines of input until the user enters an empty line
    while True:
        line = input()
        if line == "":
            break
        lines.append(line)

    # Join the lines, separated by newlines, and store the result
    user_input = "\n".join(lines)

    # Return the input
    return user_input


async def main():
    """
    Testing main function
    """
    import argparse

    print(
        """
        ChatGPT - A command-line interface to OpenAI's ChatGPT (https://chat.openai.com/chat)
        Repo: github.com/acheong08/ChatGPT
        """,
    )
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-e",
        "--email",
        help="Your OpenAI email address",
        required=False,
    )
    parser.add_argument(
        "-p",
        "--password",
        help="Your OpenAI password",
        required=False,
    )
    parser.add_argument(
        "--paid",
        help="Use the paid API",
        action="store_true",
    )
    parser.add_argument(
        "--proxy",
        help="Use a proxy",
        required=False,
        type=str,
        default=None,
    )
    parser.add_argument(
        "--insecure-auth",
        help="Use an insecure authentication method to bypass OpenAI's geo-blocking",
        action="store_true",
    )
    parser.add_argument(
        "--session_token",
        help="Alternative to email and password authentication. Use this if you have Google/Microsoft account.",
        required=False,
    )
    args = parser.parse_args()

    if (args.email is None or args.password is None) and args.session_token is None:
        print("error: " + "Please provide your email and password")
        return
    print("Logging in...")
    chatbot = Chatbot(
        args.email,
        args.password,
        paid=args.paid,
        proxy=args.proxy,
        insecure=args.insecure_auth,
        session_token=args.session_token,
    )
    print("Logged in\n")

    print("Type '!help' to show a full list of commands")
    print("Press enter twice to submit your question.\n")

    def commands(command: str) -> bool:
        if command == "!help":
            print(
                """
            !help - Show this help message
            !reset - Clear the current conversation
            !rollback <int> - Remove the latest <int> messages from the conversation
            !exit - Exit the program
            """,
            )
        elif command == "!reset":
            chatbot.conversations.remove("default")
            print("Conversation cleared")
        elif command.startswith("!rollback"):
            try:
                num = int(command.split(" ")[1])
                chatbot.conversations.rollback("default", num)
                print(f"Removed {num} messages from the conversation")
            except IndexError:
                print("Please specify the number of messages to remove")
            except ValueError:
                print("Please specify a valid number of messages to remove")
        elif command == "!exit":
            print("Exiting...")
            sys.exit(0)
        else:
            return False
        return True

    try:
        while True:
            prompt = get_input("\nYou:\n")
            if prompt.startswith("!"):
                if commands(prompt):
                    continue
            print("ChatGPT:")
            async for line in chatbot.ask(prompt=prompt):
                result = line["choices"][0]["text"].replace("<|im_end|>", "")
                print(result, end="")
                sys.stdout.flush()
            print()
    except KeyboardInterrupt:
        print("Exiting...")
        sys.exit(0)


if __name__ == "__main__":
    asyncio.run(main())
