from flask import Flask, request, jsonify, stream_with_context, Response
import requests
import json

app = Flask(__name__)

STATUS_URL = "https://duckduckgo.com/duckchat/v1/status"
CHAT_URL = "https://duckduckgo.com/duckchat/v1/chat"
STATUS_HEADERS = { "x-vqd-accept": "1" }

MODEL_MAP = {
    "gpt-4o-mini": "gpt-4o-mini",
    "claude-3-haiku": "claude-3-haiku-20240307",
    "llama": "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
    "mixtral": "mistralai/Mixtral-8x7B-Instruct-v0.1",
}

class Chat:
    def __init__(self, vqd, model):
        self.old_vqd = vqd
        self.new_vqd = vqd
        self.model = model
        self.messages = []

    def fetch(self, content):
        self.messages.append({"content": content, "role": "user"})
        payload = {
            "model": self.model,
            "messages": self.messages
        }
        response = requests.post(
            CHAT_URL,
            headers={"x-vqd-4": self.new_vqd, "Content-Type": "application/json"},
            json=payload
        )
        if not response.ok:
            response.raise_for_status()
        return response

    def fetch_full(self, content):
        response = self.fetch(content)
        text = ""
        for event in response.iter_lines():
            if event:
                message_data = json.loads(event)
                if "message" not in message_data:
                    break
                text += message_data["message"]
        
        new_vqd = response.headers.get("x-vqd-4")
        self.old_vqd = self.new_vqd
        self.new_vqd = new_vqd
        self.messages.append({"content": text, "role": "assistant"})
        return text

@app.route('/v1/models', methods=['POST', 'GET'])
def init_chat():
    model_alias = request.json.get("model")
    if model_alias not in MODEL_MAP:
        return jsonify({"error": "Invalid model alias"}), 400

    status_response = requests.get(STATUS_URL, headers=STATUS_HEADERS)
    if not status_response.ok:
        return jsonify({"error": "Failed to initialize chat"}), status_response.status_code

    vqd = status_response.headers.get("x-vqd-4")
    if not vqd:
        return jsonify({"error": "No vqd received"}), 500

    chat_instance = Chat(vqd, MODEL_MAP[model_alias])
    return jsonify({"vqd": vqd, "model": MODEL_MAP[model_alias]})

@app.route('/v1/completions', methods=['POST'])
def fetch_message():
    content = request.json.get("content")
    vqd = request.json.get("vqd")
    model = request.json.get("model")
    if not content or not vqd or not model:
        return jsonify({"error": "Missing parameters"}), 400

    chat_instance = Chat(vqd, model)
    try:
        full_message = chat_instance.fetch_full(content)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

    return jsonify({"response": full_message})

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