from flask import Flask, request, Response, jsonify
from flask_cors import CORS
import json
import re
import os
from llm.common import LlmParams, LlmPredictParams
from llm.deepinfra_api import DeepInfraApi
from llm import prompts
from dotenv import load_dotenv

# Загрузка переменных окружения из файла .env
load_dotenv()

LLM_API_URL = os.getenv("LLM_API_URL", "https://api.deepinfra.com")
LLM_API_KEY = os.getenv("DEEPINFRA_API_KEY", "")
LLM_NAME = os.getenv("LLM_NAME", "meta-llama/Llama-3.3-70B-Instruct-Turbo")

default_llm_params = LlmParams(url=LLM_API_URL,api_key=LLM_API_KEY, model=LLM_NAME, predict_params=LlmPredictParams(temperature=0.15, top_p=0.95, min_p=0.05, seed=42, repetition_penalty=1.2, presence_penalty=1.1, max_tokens=6000))
llm_api = DeepInfraApi(default_llm_params)

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


app = Flask(__name__)
CORS(app)

@app.route('/extracttable', methods=['POST'])
async def extracttable_route():
    data = request.get_json()
    text = data.get('text', '')
    
    prompt = prompts.LLM_PROMPT_EXTRACT_TABLE.format(query = text)
    response = await llm_api.predict(prompt[:150000])
    
    result = {"response": None, "error": None, "raw": response}  # По умолчанию сохраняем всю строку

    if "JSON: " not in response:
        result["error"] = "Строка не содержит 'JSON: '"
        return result

    prefix, json_str = response.split("JSON: ", 1)
    json_str = json_str.strip()

    if not json_str:
        result["error"] = "После 'JSON: ' отсутствует JSON"
        return result

    try:
        result["response"] = json.loads(json_str)
        result["raw"] = prefix.strip()  # Остаток перед "JSON: "
    except json.JSONDecodeError as e:
        result["error"] = f"Ошибка декодирования JSON: {e}"
        
    return jsonify(result)

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

@app.route('/getsummary', methods=['POST'])
async def getsummary_route():
    data = request.get_json()
    text = data.get('text', '')
    
    prompt = prompts.GET_SUMMARY.format(text=text)
    response = await llm_api.predict(prompt[:150000])
    return jsonify({'result': response})

@app.route('/cleantext', methods=['POST'])
async def cleantext_route():
    data = request.get_json()
    text = data.get('text', '')
    
    prompt = prompts.CLEAN_TEXT.format(text=text)
    response = await llm_api.predict(prompt[:150000])
    return jsonify({'result': response})

@app.route('/getfollowup', methods=['POST'])
async def getfollowup_route():
    data = request.get_json()
    text = data.get('text', '')
    
    prompt = prompts.GET_FOLLOWUP.format(text=text)
    response = await llm_api.predict(prompt[:150000])
    return jsonify({'result': response})

@app.route('/getagenda', methods=['POST'])
async def getagenda_route():
    data = request.get_json()
    text = data.get('text', '')
    
    prompt = prompts.GET_AGENDA.format(text=text)
    response = await llm_api.predict(prompt[:150000])
    return jsonify({'result': response})

@app.route('/gethighlights', methods=['POST'])
async def gethighlights_route():
    data = request.get_json()
    text = data.get('text', '')
    
    prompt = prompts.GET_HIGHLIGHTS.format(text=text)
    response = await llm_api.predict(prompt[:150000])
    return jsonify({'result': response})


@app.route('/getprojectinfo', methods=['POST'])
async def getprojectinfo_route():
    data = request.get_json()
    text = data.get('text', '')

    main_prompts = []
    main_prompts.append(prompts.GET_PROJECT_INFO_NAMES.format(text=text))
    main_prompts.append(prompts.GET_PROJECT_INFO_AGENDA.format(text=text))
    
    main_info =''
    for i in main_prompts:
        result = await llm_api.predict(i[:150000])
        if result is None:
            return jsonify({'error': 'Сервер LLM временно недоступен. Попробуйте повторить запрос через несколько минут.'})
        main_info += '\n\n'+result+'\n\n'
    
    final = main_info
    final = final.replace("Конец ответа", "")
    final = final.replace('</s>', '')
    final = final.strip()
    return jsonify({'result': final})


@app.route('/getprojectlist', methods=['POST'])
async def getprojectlist_route():
    data = request.get_json()
    text = data.get('text', '')

    main_prompts = []
    main_prompts.append(prompts.GET_PROJECT_INFO_NAMES.format(text=text))
    main_prompts.append(prompts.GET_PROJECT_INFO_AGENDA.format(text=text))
    
    main_info =''
    for i in main_prompts:
        result = await llm_api.predict(i[:150000])
        if result is None:
            return jsonify({'error': 'Сервер LLM временно недоступен. Попробуйте повторить запрос через несколько минут.'})
        main_info += '\n\n'+result+'\n\n'

    proj_prompt = []
    proj_prompt.append(prompts.GET_PROJECT_LIST.format(text=text))
    
    list_of_projects =''
    for i in proj_prompt:
        result = await llm_api.predict(i[:150000])
        if result is None:
            return jsonify({'error': 'Сервер LLM временно недоступен. Попробуйте повторить запрос через несколько минут.'})
        list_of_projects += result
    
    delimiter = 'Проект '
    proj  = [delimiter+x for x in list_of_projects.split(delimiter) if x]
    proj = proj[1:]
    
    projects = []
    for i in proj:
        a = i.replace("Проект №", "")
        a = a.replace("Конец ответа", "")
        a = a.replace("данный проект", "")  ###убираю слово "проект", чтобы модель не опиралась на него,
        a = a.replace("проект ", "")             # при ответе на вопрос, проект это или нет
        a = a.replace('\n', ' ')
        a = a.replace('</s>', ' ')
        a = a.strip()
        projects.append(a)
        
    check_prompts = []

    checking = prompts.GET_PROJECT_LIST_CHECK_PROJECT.format(text=text, projects=projects)
    check_prompts.append(checking)

    real_projects = ''
    for i in check_prompts:
        result = await llm_api.predict(i[:150000])
        if result is None:
            return jsonify({'error': 'Сервер LLM временно недоступен. Попробуйте повторить запрос через несколько минут.'})
        real_projects += result
    
    real_projects_list = re.findall(r'Да:\s*(.*?)\s*(?:\n\n|$)', real_projects)
    return jsonify({'result': real_projects_list})

@app.route('/getprojectdetails', methods=['POST'])
async def getinfobyproject_route():
    data = request.get_json()
    text = data.get('text', '')
    real_projects_list = data.get('projects', {})
    
    project_prompts = {}
    if real_projects_list:
        for i in real_projects_list:
            if not i or i.strip() == "":
                continue
            
            prompt_aim = prompts.GET_PROJECT_DETAILS_AIM.format(text=text, project=i)
            gk = prompts.GET_PROJECT_DETAILS_VALUE.format(text=text, project=i)
            budget = prompts.GET_PROJECT_DETAILS_BUDGET.format(text=text, project=i)
            ec_ef = prompts.GET_PROJECT_DETAILS_ECO_EFFECT.format(text=text, project=i)
            deadline = prompts.GET_PROJECT_DETAILS_DEADLINE.format(text=text, project=i)
            new_plan = prompts.GET_PROJECT_DETAILS_NEW_PLAN.format(text=text, project=i)
            conclusion = prompts.GET_PROJECT_DETAILS_CONCLUSION.format(text=text, project=i)
            
            p = [prompt_aim, gk, budget, ec_ef, deadline, new_plan, conclusion]
            project_prompts[i] = {}
            project_prompts[i]['prompts'] = p
        
    elif not real_projects_list:
        return jsonify({'error': 'Проекты не выбраны'})
    
    final = {}
    for project_name, project in project_prompts.items():
        for prompt in project['prompts']:
            result = await llm_api.predict(prompt[:150000])
            if result is not None:
                final[project_name] = final.get(project_name, '') + '\n\n'+result + '\n\n'
        final[project_name] = final[project_name].replace("Конец ответа", "")
        final[project_name] = final[project_name].replace('</s>', '')
        final[project_name] = final[project_name].strip()

    return jsonify({'result': final})


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