import os
import streamlit as st
import requests
import json
import math

def get_completion_from_openai(prompt, max_tokens = None):
    url = os.getenv('OPENAI_COMPLETION_URL')

    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer " + st.secrets["OPENAI_TOKEN"],
    }

    response = requests.post(url,
                                 json={    
                                     "model": os.getenv('OPENAI_MODEL'),
                                     "max_tokens": max_tokens,
                                     "messages": [
                                      {
                                        "role": "user",
                                        "content": prompt
                                      }
                                    ]
                                }, 
                                 headers=headers, 
                                 stream=False,
                                 )
    try:
        return response.json()['choices'][0]['message']['content']
    except:
        print(response.json())
        return "Произошла ошибка"

def process_transcribation_with_assistant(prompt, transcript):
    baseUrl = os.getenv('OPENAI_BASE_URL')

    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer " + st.secrets["OPENAI_TOKEN"],
        "OpenAI-Beta": "assistants=v2",
    }

    number_of_runs = math.ceil(len(transcript) / (4 * 4096))
    output_text = ''

    thread_response = requests.post(baseUrl + '/threads', json={}, headers=headers)
    thread_id = thread_response.json()['id']

    message_response = requests.post(baseUrl + '/threads/' + thread_id + '/messages', 
                                     json={"role" : "user", "content": prompt + transcript}, 
                                     headers=headers)

    run_response = requests.post(baseUrl + '/threads/' + thread_id + '/runs',
                             json={    
                                 "assistant_id": st.secrets["OPENAI_ASSISTANT_ID"],
                                 "stream": True
                                 }, 
                             headers=headers, 
                             stream=True)

    st.write("Результат обработки:")
    text_container = st.empty()
    output_text = ''
    event_name = ""
    for line in run_response.iter_lines(decode_unicode=True):
        if line:
            if line.startswith("event:"):
                event_name = line.split(":")[1].strip()
                if event_name == 'done':
                    break
            elif line.startswith("data:") and event_name == 'thread.message.delta':
                event_data = json.loads(line.split(":", 1)[1].strip())
                output_text += event_data['delta']['content'][0]['text']['value']
                text_container.text(output_text)

    for i in range(number_of_runs - 1):
        message_response = requests.post(baseUrl + '/threads/' + thread_id + '/messages', 
                                          json={"role" : "user", "content": "Продолжай работать на текущей задачей"}, 
                                          headers=headers)

        run_response = requests.post(baseUrl + '/threads/' + thread_id + '/runs',
                             json={    
                                 "assistant_id": st.secrets["OPENAI_ASSISTANT_ID"],
                                 "stream": True
                                 }, 
                             headers=headers, 
                             stream=True)
        event_name = ""
        response = ''
        for line in run_response.iter_lines(decode_unicode=True):
            if line:
                if line.startswith("event:"):
                    event_name = line.split(":")[1].strip()
                    if event_name == 'done':
                        break
                elif line.startswith("data:") and event_name == 'thread.message.delta':
                    event_data = json.loads(line.split(":", 1)[1].strip())
                    output_text += event_data['delta']['content'][0]['text']['value']
                    text_container.text(output_text)

    return output_text