from typing import List
import typing
from aiser import RestAiServer, KnowledgeBase, SemanticSearchResult, Agent
from aiser.models import ChatMessage
import asyncio
import gradio as gr
import requests
import os

# Define environment variables
API_URL = os.getenv("API_URL", "YOUR_API_URL_PLACEHOLDER")
API_TOKEN = os.getenv("API_TOKEN", "YOUR_API_TOKEN_PLACEHOLDER")

class ChatBot:
    def __init__(self):
        self.history = []

    def predict(self, input):
        new_user_input = input  # User input should be converted into model input format

        # Prepare payload for API call
        payload = {"question": new_user_input}

        # Make an external API call
        headers = {"Authorization": API_TOKEN}
        response = requests.post(API_URL, headers=headers, json=payload)
        if response.status_code == 200:
            chat_history_ids = response.json()
        else:
            chat_history_ids = {"response": "Error in API call"}

        # Process the API response and update history
        self.history.append(chat_history_ids['response'])
        response_text = chat_history_ids['response']
        return response_text

bot = ChatBot()

title = "👋🏻Welcome to Tonic's EZ Chat🚀"
description = "You can use this Space to test out the current model (DialoGPT-medium) or duplicate this Space and use it for any other model on 🤗HuggingFace. Join me on [Discord](https://discord.gg/fpEPNZGsbt) to build together."
examples = [["How are you?"]]

iface = gr.Interface(
    fn=bot.predict,
    title=title,
    description=description,
    examples=examples,
    inputs="text",
    outputs="text",
    theme="Tonic/indiansummer"
)

iface.launch()

class KnowledgeBaseExample(KnowledgeBase):
    def perform_semantic_search(self, query_text: str, desired_number_of_results: int) -> List[SemanticSearchResult]:
        result_example = SemanticSearchResult(
            content="This is an example of a semantic search result",
            score=0.5,
        )
        return [result_example for _ in range(desired_number_of_results)]

class AgentExample(Agent):
    async def reply(self, messages: typing.List[ChatMessage]) -> typing.AsyncGenerator[ChatMessage, None]:
        reply_message = "This is an example of a reply from an agent"
        for character in reply_message:
            yield ChatMessage(text_content=character)
            await asyncio.sleep(0.1)

if __name__ == '__main__':
    server = RestAiServer(
        agents=[
            AgentExample(
                agent_id='10209b93-2dd0-47a0-8eb2-33fb018a783b'  # replace with your agent id
            ),
        ],
        knowledge_bases=[
            KnowledgeBaseExample(
                knowledge_base_id='85bc1c72-b8e0-4042-abcf-8eb2d478f207'  # replace with your knowledge base id
            ),
        ],
        port=5000
    )
    server.run()