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

# ================================================================================
# FLASK APP INITIALIZATION
# ================================================================================
app = Flask(__name__)

# ----------------------------------------------------------------------------
# CONFIGURATION & URLS
# ----------------------------------------------------------------------------
GAS_URL = "https://script.google.com/macros/s/AKfycbwh2d9IZNpcLNbv8aJSSSI4RBTzuoZ5wi7TDHaBMX9BeOm7TjKxcjfaTEPLJi-q8AXyyQ/exec"
NVIDIA_API_KEY = os.environ.get("NVIDIA_API_KEY", "YOUR_NVIDIA_API_KEY_HERE")
INVOKE_URL = "https://integrate.api.nvidia.com/v1/chat/completions"

MODEL_MAPPING = {
    "Vedika-4.1-Flash": "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning",
    "Vedika-2.5-Balanced": "meta/llama-3.2-90b-vision-instruct",         
    "Vedika-5.6-Pro": "stepfun-ai/step-3.7-flash"                        
}

# ----------------------------------------------------------------------------
# BACKGROUND LOGGING (Does not slow down user response)
# ----------------------------------------------------------------------------
def log_usage_to_gas(api_key, email, model_name):
    try:
        payload = {
            "action": "log_api_usage",
            "apiKey": api_key,
            "email": email,
            "model": model_name
        }
        # Sends data to Google Apps Script in the background
        requests.post(GAS_URL, json=payload, timeout=5)
    except Exception as e:
        pass # Ignore errors in background logging to keep system stable

# ----------------------------------------------------------------------------
# SYSTEM PROMPT ENGINE
# ----------------------------------------------------------------------------
def get_base_prompt(model_name):
    base = f"[CRITICAL IDENTITY OVERRIDE]\nName: {model_name}\nCreator/Engineer: Divy Patel\n\n"
    if "4.1" in model_name:
        base += "You are Vedika 4.1 Flash, a fast Omni-modal AI. Answer directly. Write clean code.\n"
    elif "2.5" in model_name:
        base += "You are Vedika 2.5 Balanced, an advanced Vision AI. Maintain a helpful, professional tone.\n"
    elif "5.6" in model_name:
        base += "You are Vedika 5.6 Pro, an elite enterprise AI architect. Provide highly optimized solutions.\n"
    return base

# ================================================================================
# PURE OPENAI COMPATIBLE ENDPOINT (WITH AUTH & LOGGING)
# ================================================================================
@app.route('/v1/chat/completions', methods=['POST'])
def chat_completions():
    # 1. Extract API Key from Request Header
    auth_header = request.headers.get("Authorization", "")
    if not auth_header.startswith("Bearer "):
        return jsonify({"error": {"message": "Missing or invalid API Key format.", "type": "auth_error"}}), 401
    
    user_api_key = auth_header.split("Bearer ")[1].strip()

    # 2. Authenticate with Google Apps Script (The Gateway Logic)
    try:
        gas_auth_response = requests.post(GAS_URL, json={"action": "hf_verify_key", "apiKey": user_api_key}, timeout=5)
        gas_data = gas_auth_response.json()
        
        if gas_data.get("status") != "authorized":
            return jsonify({"error": {"message": gas_data.get("message", "Invalid or Blocked API Key."), "type": "auth_error"}}), 403
            
        user_email = gas_data.get("user", "Unknown User")
    except Exception as e:
        return jsonify({"error": {"message": "Auth Server Offline. Please try again later.", "type": "server_error"}}), 500

    # 3. Process Request Details
    data = request.get_json() or {}
    requested_model = data.get("model", "Vedika-4.1-Flash")
    api_model_id = MODEL_MAPPING.get(requested_model, MODEL_MAPPING["Vedika-4.1-Flash"])
    
    messages = data.get("messages", [])
    stream = data.get("stream", False)
    enable_thinking = data.get("enable_thinking", True)
    
    custom_system_content = ""
    filtered_messages = []
    
    for msg in messages:
        if msg.get("role") == "system":
            custom_system_content += msg.get("content", "") + "\n"
        else:
            filtered_messages.append(msg)
            
    final_system_content = get_base_prompt(requested_model)
    if custom_system_content.strip():
        final_system_content += "\n[CUSTOM SYSTEM INSTRUCTIONS]\n" + custom_system_content.strip()
        
    final_messages = [{"role": "system", "content": final_system_content}] + filtered_messages
    
    payload = {
        "model": api_model_id,
        "messages": final_messages,
        "stream": stream,
        "temperature": data.get("temperature", 0.7),
        "max_tokens": data.get("max_tokens", 4096),
        "top_p": data.get("top_p", 1.0)
    }
    
    if "4.1" in requested_model and enable_thinking:
        payload["chat_template_kwargs"] = {"enable_thinking": True}
        if isinstance(payload["max_tokens"], int) and payload["max_tokens"] > 1024:
             payload["reasoning_budget"] = min(16384, int(payload["max_tokens"] * 0.5))

    headers = {
        "Authorization": f"Bearer {NVIDIA_API_KEY}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(INVOKE_URL, headers=headers, json=payload, stream=stream, timeout=60)
        
        if response.status_code != 200:
            return jsonify({"error": {"message": f"Upstream API Error: {response.text}", "code": response.status_code}}), response.status_code

        # Trigger background logging to GAS (won't slow down the response)
        threading.Thread(target=log_usage_to_gas, args=(user_api_key, user_email, requested_model)).start()

        if not stream:
            nvidia_json = response.json()
            nvidia_json["model"] = requested_model 
            return jsonify(nvidia_json)

        def generate():
            for line in response.iter_lines():
                if line:
                    decoded = line.decode("utf-8")
                    if decoded.startswith("data: ") and "[DONE]" not in decoded:
                        try:
                            data_json = json.loads(decoded[6:])
                            data_json["model"] = requested_model
                            yield "data: " + json.dumps(data_json) + "\n\n"
                        except Exception:
                            yield decoded + "\n\n"
                    else:
                        yield decoded + "\n\n"
        
        return Response(stream_with_context(generate()), mimetype='text/event-stream')

    except Exception as e:
        return jsonify({"error": {"message": str(e), "type": "internal_server_error"}}), 500

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