from flask import Flask, request, Response, jsonify
from huggingface_hub import InferenceClient
from flask_cors import CORS
import json
import requests

client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

summary_prompt = '<s>[INST]Ты ассистент. Отвечаешь на русском языке. Сформируй краткое изложение следующего текста: {}[/INST]'

def format_prompt(message, history):
  prompt = "<s>"
  for user_prompt, bot_response in history:
    prompt += f"[INST] {user_prompt} [/INST]"
    prompt += f" {bot_response}</s> "
  prompt += f"[INST] {message} [/INST]"
  return prompt

def split_text(text):
    max_chars = 3500
    sentences = text.split('.')
    lines = []
    for sentence in sentences:
        lines.extend(sentence.split('\n'))

    result = []
    current_chunk = ''
    for line in lines:
        if len(current_chunk) + len(line) < max_chars:
            current_chunk += line + '.'
        else:
            result.append(current_chunk.strip())
            current_chunk = line + '.'
    if current_chunk:
        result.append(current_chunk.strip())
    return result

def call_api(prompt_text):
    url = "https://muryshev-mixtral-api-local.hf.space/completion"
    payload = {"prompt": prompt_text, "temperature": 0, "seed": 42, "repeat_penalty": 1, "top_p": 0.95, "stream": False, "n_predict": 2000}
    
    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()  # Raise an exception for 4xx or 5xx status codes
        result = response.json().get('content', '') # Extract the text result from the JSON response
        print(result)
        return result
    except requests.exceptions.RequestException as e:
        print("Error:", e)
        return None
    
def generate(
    prompt, history=[], temperature=0, max_new_tokens=2000, top_p=0.95, repetition_penalty=1.0,
):
    return call_api(prompt)
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )

    #formatted_prompt = format_prompt(prompt, history)

    #stream = client.text_generation(prompt, **generate_kwargs, stream=True, details=False, return_full_text=False)
    response = client.text_generation(prompt, **generate_kwargs, stream=False, details=False, return_full_text=False)
    print(response)
    return response
    #output = ""

    #for response in stream:
    #    yield response.token.text.encode('utf-8')

app = Flask(__name__)
CORS(app)

@app.route('/health', methods=['GET'])
def health():
    return jsonify({"status": "ok"})

@app.route('/completion', methods=['POST'])
def completion_route():
    data = request.get_json()
    prompt = data.get('prompt', '')
    #truncated_prompt = prompt[:32768]
    return Response(generate(prompt[:52768]), content_type='text/plain; charset=utf-8', status=200, direct_passthrough=True)

@app.route('/getsummary', methods=['POST'])
def getsummary_route():
    data = request.get_json()
    text = data.get('text', '')
    # pages = split_text(text)
    # result = ''
    # for page in pages:
    #     summary_prompt = f'''<s>[INST]Исправь в тексте, сделанным роботом при транскрибации аудиозаписи совещания, плохо-читаемые предложения. Проверь каждое предложение на согласованность падежей, окончаний, на наличие и правильное использование предлогов, в реплике каждого спикера исправь пунктуацию, чтобы смысл и был четкий и понятный, также удали дискурсивные слова, не несущие смысла, сделай текст аккуратным и легко читабельным и пиши только на русском языке: {page}[/INST]'''
    #     response = generate(summary_prompt[:52000])
    #     result = result + '\n'+response
    
    summary_prompt = f'''<s>[INST]Исправь в тексте, сделанным роботом при транскрибации аудиозаписи совещания, плохо-читаемые предложения. Проверь каждое предложение на согласованность падежей, окончаний, на наличие и правильное использование предлогов, в реплике каждого спикера исправь пунктуацию, чтобы смысл и был четкий и понятный, также удали дискурсивные слова, не несущие смысла, сделай текст аккуратным и легко читабельным и пиши только на русском языке: {text}[/INST]'''
    result = generate(summary_prompt[:52000])
    return jsonify({'result': result})

@app.route('/cleantext', methods=['POST'])
def cleantext_route():
    data = request.get_json()
    text = data.get('text', '')
    
    summary_prompt = f'''<s>[INST]Ты мой помощник. Ты отвечаешь только на русском языке. Сформируй краткое изложение следующего текста: {text}[/INST]'''
    response = generate(summary_prompt[:52000])
    
    return jsonify({'result': response})

@app.route('/getfollowup', methods=['POST'])
def getfollowup_route():
    data = request.get_json()
    text = data.get('text', '')
    
    summary_prompt = f'''<s>[INST]Ты мой помощник. Ты отвечаешь только на русском языке. Выпиши из текста нумерованным списком обсуждаемые темы и у каждой темы сделай подпунктами обсуждаемые вопросы. Текст: {text}[/INST]'''
    response = generate(summary_prompt[:52000])
    
    return jsonify({'result': response})

@app.route('/getagenda', methods=['POST'])
def getagenda_route():
    data = request.get_json()
    text = data.get('text', '')
    
    summary_prompt = f'''<s>[INST]Ты мой помощник. Ты отвечаешь только на русском языке. Выпиши из текста в виде списка какие задачи были поставлены для конкретных исполнителей с указанием сроков, если это возможно определить. Если задача есть в списке, то не добавляй идентичную. Текст: {text}[/INST]'''
    response = generate(summary_prompt[:52000])
    
    return jsonify({'result': response})

@app.route('/gethighlights', methods=['POST'])
def gethighlights_route():
    data = request.get_json()
    text = data.get('text', '')
    
    summary_prompt = f'''<s>[INST]Ты мой помощник. Ты отвечаешь только на русском языке. Сформируй на основе текста два нумерованных списка: 1. какие проблемы были озвучены в тексте 2. какие предложения были сформулированы. Текст: {text}[/INST]'''
    response = generate(summary_prompt[:52000])
    
    return jsonify({'result': response})

if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0', port=7860)