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

app = Flask(__name__)
CORS(app)

def get_random_api_key():
    keys = os.getenv("KEYS", "").split(",")
    if keys and keys[0]:
        return random.choice(keys).strip()
    else:
        raise ValueError("API keys not found. Please set the KEYS environment variable.")

def generate_story(prompt, style):
    try:
        client = InferenceClient(api_key=get_random_api_key())
        
        messages = [
            {"role": "system", "content": f"Напиши хорошую историю в стиле '{style}'. Подробную, понятную, человечную (с душой), уникальную. Не обязательно делать концовку, можно только начало длинной истории. Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь."},
            {"role": "user", "content": prompt}
        ]
        
        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=True)
        for chunk in completion:
            if chunk.choices:
                yield chunk.choices[0].delta.content or ""
    except Exception as e:
        yield f"Ошибка генерации: {e}"

def edit_story(original_story, edited_prompt):
    if not original_story:
        yield f"Сначала сгенерируйте историю!"
        return
    if not edited_prompt:
        yield f"Для начала введите что изменить."
        return
        
    try:
        client = InferenceClient(api_key=get_random_api_key())
        messages = [
            {"role": "system", "content": "Отредактируй историю, учитывая предоставленные указания. Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь."},
            {"role": "user", "content": edited_prompt},
            {"role": "user", "content": f"""```История
{original_story}
```"""}
        ]
        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=155000, stream=True)
        for chunk in completion:
            if chunk.choices:
                yield chunk.choices[0].delta.content or ""
    except Exception as e:
        yield f"Ошибка редактирования: {e}"

def next_story_func(original_story, next_prompt, continuation_type="Продолжение"):
    if not original_story:
        yield f"Сначала сгенерируйте историю!"
        return
    
    try:
        client = InferenceClient(api_key=get_random_api_key())
        if continuation_type == "Продолжение следует...":
            continuation_prompt = "Продолжи, подведя к тому, что продолжение следует."
        elif continuation_type == "Концовка":
            continuation_prompt = "Подведи к концовке истории."
        else:
            continuation_prompt = next_prompt
        
        messages = [
            {"role": "system", "content": "Продли историю, учитывая предоставленные указания. Продливай В ТОЧНОСТИ С КОНЦА, прям с того же символа, слова, предложения. (В начале добавляй новые строки/пробел если надо для отступа). Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь."},
            {"role": "user", "content": continuation_prompt},
            {"role": "assistant", "content": f"""```История
{original_story}
```"""}
        ]
        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=True)
        for chunk in completion:
            if chunk.choices:
                yield chunk.choices[0].delta.content or ""
    except Exception as e:
        yield f"Ошибка продления: {e}"

def stream_sse(generator):
    def event_stream():
       try:
           for chunk in generator:
              yield f"data: {json.dumps({'message': chunk})}\n\n"
       except Exception as e:
           yield f"data: {json.dumps({'error': True, 'message': str(e)})}\n\n"
    return Response(stream_with_context(event_stream()), content_type='text/event-stream')

@app.route('/generate', methods=['POST'])
def api_generate_story():
    data = request.get_json()
    if not data or 'input' not in data:
        return jsonify({"error": True, "message": "Missing 'input' in request body"}), 400
    
    prompt = data['input']
    style = data.get('style', 'Приключенческая')
    
    story_generator = generate_story(prompt, style)
    return stream_sse(story_generator)

@app.route('/edit', methods=['POST'])
def api_edit_story():
    data = request.get_json()
    if not data or 'input' not in data or 'original' not in data:
        return jsonify({"error": True, "message": "Missing 'input' or 'original' in request body"}), 400
    
    original_story = data['original']
    edited_prompt = data['input']
    
    edited_story_generator = edit_story(original_story, edited_prompt)
    return stream_sse(edited_story_generator)

@app.route('/continue', methods=['POST'])
def api_continue_story():
    data = request.get_json()
    if not data or 'input' not in data or 'original' not in data:
        return jsonify({"error": True, "message": "Missing 'input' or 'original' in request body"}), 400
    
    original_story = data['original']
    next_prompt = data['input']
    continuation_type = data.get('type', 'Продолжение')
    
    next_story_generator = next_story_func(original_story, next_prompt, continuation_type)
    return stream_sse(next_story_generator)


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