import gradio as gr

from utils.chatgpt import ChatGPTAPI
from utils.read_pdf import read_pdf
from utils.read_web import read_web
from utils.truncate import truncate_string


def file2str(filepath: str) -> str:
    if not filepath:
        return ''
    if filepath.endswith('.pdf'):
        content_list = read_pdf(filepath)
        text = '\n'.join(content_list)
    elif filepath.endswith('.txt'):
        with open(filepath, 'r') as f:
            text = f.readlines()
    else:
        raise Exception('File type not supported')
    text = truncate_string(text, max_length=1024)
    return text


def process(api_key: str = '', prompt: str = '', file=None, url='') -> str:
    chatgpt = ChatGPTAPI(api_key, max_input_length=1024)

    file_text = file2str(file.name) if file else ''
    web_txt = read_web(url)
    web_txt = truncate_string(web_txt, max_length=1024)

    content = prompt + '\n' + file_text + '\n' + web_txt
    response = chatgpt(content)
    return response


prompt_input = gr.components.Textbox(
    value='用中文总结下面的文章',
    lines=2,
    type="text"
)

app = gr.Interface(
    fn=process,
    inputs=["text", prompt_input, "file", "text"],
    outputs="text"
)
app.launch()