import os
import random
import requests
from huggingface_hub import InferenceClient
from flask import Flask, request, jsonify
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, stream=False):
    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=stream)
        if stream:
            story = ""
            for chunk in completion:
                if chunk.choices:
                    story += chunk.choices[0].delta.content or ""
                    yield story
        else:
            story = completion.choices[0].message.content
            yield story
    except Exception as e:
        yield f"Ошибка генерации: {e}"

def edit_story(original_story, edited_prompt, stream=False):
    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": "assistant", "content": original_story}
        ]
        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=155000, stream=stream)
        if stream:
            edited_story = ""
            for chunk in completion:
                if chunk.choices:
                    edited_story += chunk.choices[0].delta.content or ""
                    yield edited_story
        else:
            edited_story = completion.choices[0].message.content
            yield edited_story
    except Exception as e:
        yield f"Ошибка редактирования: {e}"

def next_story_func(original_story, next_prompt, continuation_type="Продолжение", stream=False):
    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": original_story}
        ]
        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=stream)
        if stream:
            next_story = ""
            for chunk in completion:
                if chunk.choices:
                    next_story += chunk.choices[0].delta.content or ""
                    yield next_story
        else:
            next_story = completion.choices[0].message.content
            yield next_story
    except Exception as e:
        yield f"Ошибка продления: {e}"

@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', 'Приключенческая')  # Default style
    stream = data.get('stream', False)
    
    try:
        story_generator = generate_story(prompt, style, stream)
        if stream:
            return jsonify({"error": False, "message": "".join(story_generator)})
        else:
            return jsonify({"error": False, "message": next(story_generator)})
    except Exception as e:
        return jsonify({"error": True, "message": str(e)}), 500

@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']
    stream = data.get('stream', False)
    
    try:
        edited_story_generator = edit_story(original_story, edited_prompt, stream)
        if stream:
           return jsonify({"error": False, "message": "".join(edited_story_generator)})
        else:
            return jsonify({"error": False, "message": next(edited_story_generator)})
    except Exception as e:
        return jsonify({"error": True, "message": str(e)}), 500

@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', 'Продолжение') # Default continuation type
    stream = data.get('stream', False)
    
    try:
        next_story_generator = next_story_func(original_story, next_prompt, continuation_type, stream)
        if stream:
           return jsonify({"error": False, "message": "".join(next_story_generator)})
        else:
           return jsonify({"error": False, "message": next(next_story_generator)})
    except Exception as e:
        return jsonify({"error": True, "message": str(e)}), 500

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