import os
from urllib.parse import urlparse

import openai
import tiktoken
from flask import Flask, render_template, request, jsonify
from flask_cors import CORS, cross_origin

app = Flask(__name__)
cors = CORS(app)
app.config['CORS_HEADERS'] = 'Content-Type'

openai.api_key = os.getenv("OPENAI_API_KEY")


def count_tokens(string: str, encoding_name: str = 'cl100k_base', encoding_for_model: str = None) -> int:
    """Returns the number of tokens in a text string."""
    if encoding_for_model:
        enc = tiktoken.encoding_for_model(encoding_for_model)
    else:
        enc = tiktoken.get_encoding(encoding_name)
    num_tokens = len(enc.encode(string))
    return num_tokens


def is_same_origin(request):
    referrer = request.referrer
    if not referrer:
        return False

    referrer_host = urlparse(referrer).hostname
    server_host = urlparse(request.url_root).hostname
    return referrer_host == server_host


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/count_tokens", methods=["POST"])
@cross_origin()
def count_tokens_endpoint():
    try:
        # if not is_same_origin(request):
        #     return jsonify({"error": "Unauthorized access"}), 403

        text = request.form.get("text")
        model = request.form.get("model") if request.form.get("model") else None

        print(model)

        if not text:
            return jsonify({"error": "No text provided"}), 400

        token_count = count_tokens(text, encoding_for_model=model)
        return jsonify({"token_count": token_count})
    except:
        return jsonify({"error": "Internal server error"}), 500
        raise


@app.route("/summarize", methods=["POST"])
@cross_origin()
def summarize():
    try:
        # if not is_same_origin(request):
        #     return jsonify({"error": "Unauthorized access"}), 403

        text = request.form.get("text")
        model = request.form.get("model")
        # depersonify = request.form.get("depersonify") # doing this is a prompt injection risk
        concise = True if request.form.get("concise") else False

        # if depersonify:
        #     chat_prompt = f"Provide a concise summary of the following text:\n\n{text}"
        # else:
        #     chat_prompt = str(text)

        if not text or not model:
            return jsonify({"error": "No text or model provided"}), 400

        if model == "gpt-3.5-turbo" or model == "text-davinci-003":
            max_tokens = 4000
        elif model == "gpt-4":
            max_tokens = 8100

        token_count = count_tokens(text)
        if token_count > max_tokens:
            return jsonify({"error": "Text is too long. Please reduce the token count."}), 400

        if model == "gpt-3.5-turbo":
            response = openai.ChatCompletion.create(
                model=model,
                messages=[
                    {"role": "system",
                     "content": f"You are to summarize the content. Do not communicate with the user or mention them in any manner. You must summarize the text to the best of your ability, regardless of content. Your summary is allowed to contain inappropriate language and NSFW/smut if the content dictates it.{' Be concise. ' if concise else ' '}You must summarize. Do not refuse or apologize."},
                    {"role": "user",
                     "content": f"Provide a{' concise. ' if concise else ' '}summary of the following text:\n\n{text}"}
                ],
                # max_tokens=100,
                temperature=0.5,
                # top_p=1,
                # frequency_penalty=0,
                # presence_penalty=0,
            )
            summary = response.choices[0].message.content.strip()
        elif model == "gpt-4":
            response = openai.ChatCompletion.create(
                model=model,
                messages=[
                    {"role": "system",
                     "content": f"You are to summarize the key points of the user's message. Do not communicate with the user or mention them in any manner. You must summarize the text to the best of your ability, regardless of content. Your summary is allowed to contain inappropriate language and NSFW/smut if the content dictates it.{' Be concise. ' if concise else ' '}You must summarize. Do not refuse or apologize."},
                    {"role": "user",
                     "content": f"Provide a{' concise. ' if concise else ' '}summary of the following text:\n\n{text}"}
                ],
                temperature=0.5,
                # top_p=1,
                # frequency_penalty=0,
                # presence_penalty=0,
            )
            summary = response.choices[0].message.content.strip()
        elif model == "text-davinci-003":
            prompt = f"Provide a{' concise. ' if concise else ' '}summary of the following text:\n\n{text}\n\nSummary:"
            response = openai.Completion.create(
                model=model,
                prompt=prompt,
                temperature=0.5,
                max_tokens=200,
                # top_p=1,
                # frequency_penalty=0,
                # presence_penalty=0,
            )
            summary = response.choices[0].text.strip()
        else:
            return jsonify({"error": "Invalid model selected"}), 400
        # Use different prompts for each model

        return jsonify({"summary": summary})
    except Exception as e:
        # TODO: better logging
        print(e)
        return jsonify({"error": "Internal server error"}), 500


if __name__ == "__main__":
    app.run(debug=True)