from flask import Flask, request, jsonify
import sqlite3
import uuid
import subprocess
import time

app = Flask(__name__)
DB_PATH = "/var/www/deepseek_api/chat_contexts.db"

@app.route('/initiate', methods=['POST'])
def initiate_chat():
    context_id = str(uuid.uuid4())
    lifespan = request.json.get('lifespan', 3600)  # Default to 1 hour
    conn = sqlite3.connect(DB_PATH)
    conn.execute("INSERT INTO contexts (id, messages, lifespan) VALUES (?, ?, ?)", (context_id, "[]", lifespan))
    conn.commit()
    conn.close()
    return jsonify({"context_id": context_id, "lifespan": lifespan})

@app.route('/chat/<context_id>', methods=['POST'])
def continue_chat(context_id):
    user_message = request.json.get('message')
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute("SELECT messages, created_at, lifespan FROM contexts WHERE id = ?", (context_id,))
    result = cursor.fetchone()
    if not result:
        return jsonify({"error": "Invalid context_id"}), 404
    messages, created_at, lifespan = result
    # Check if context is expired
    created_timestamp = time.mktime(time.strptime(created_at, "%Y-%m-%d %H:%M:%S"))
    if time.time() - created_timestamp > lifespan:
        cursor.execute("DELETE FROM contexts WHERE id = ?", (context_id,))
        conn.commit()
        conn.close()
        return jsonify({"error": "Context expired"}), 410
    messages = eval(messages)
    messages.append({"user": user_message})
    # Generate response using DeepSeek in Ollama
    response = subprocess.run(["ollama", "run", "deepseek", user_message], capture_output=True, text=True).stdout.strip()
    messages.append({"deepseek": response})
    cursor.execute("UPDATE contexts SET messages = ? WHERE id = ?", (str(messages), context_id))
    conn.commit()
    conn.close()
    return jsonify({"response": response, "context_id": context_id})

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