import os
import logging
import requests
import json
from flask import Flask, request, jsonify, Response, stream_with_context, send_file
from flask_cors import CORS
from openai import OpenAI

# ---------------------------------------------------------
# लॉगिंग सेटअप
# ---------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - [Vedika Space] - %(message)s'
)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

# ---------------------------------------------------------
# Secrets और Configuration
# ---------------------------------------------------------
API_KEY = os.getenv("API_KEY")
MODEL_ID = os.getenv("MODEL_ID")
GAS_BACKEND_URL = os.getenv("GAS_BACKEND_URL")

# OpenAI क्लाइंट इनिशियलाइज़ेशन (NVIDIA बेस यूआरएल के साथ)
client = OpenAI(
  base_url="https://integrate.api.nvidia.com/v1",
  api_key=API_KEY
)

# ---------------------------------------------------------
# सिक्योरिटी: Google Apps Script टोकन वैलिडेशन
# ---------------------------------------------------------
def validate_token_with_gas(token):
    if not token:
        return False
    try:
        response = requests.post(GAS_BACKEND_URL, json={"action": "validate", "token": token}, timeout=5)
        if response.status_code == 200:
            return response.json().get("valid") == True
        return False
    except requests.exceptions.RequestException as e:
        logger.error(f"GAS validation error: {e}")
        return False

# ---------------------------------------------------------
# फ्रंटएंड रेंडरिंग
# ---------------------------------------------------------
@app.route('/')
def index():
    return send_file('index.html')

# ---------------------------------------------------------
# बैकएंड API (अपडेटेड OpenAI SDK के साथ)
# ---------------------------------------------------------
@app.route('/chat', methods=['POST', 'OPTIONS'])
def chat_endpoint():
    # CORS प्री-फ्लाइट रिक्वेस्ट
    if request.method == 'OPTIONS':
        return Response(status=204)

    # ऑथेंटिकेशन
    auth_header = request.headers.get("Authorization")
    if not auth_header or not auth_header.startswith("Bearer "):
        return jsonify({"error": "Unauthorized: Missing Token"}), 401

    token = auth_header.split("Bearer ")[1].strip()

    if not validate_token_with_gas(token):
        return jsonify({"error": "Forbidden: Invalid Token"}), 403

    try:
        data = request.json
        mode = data.get('mode', 'FLASH')
        messages = data.get('messages', [])
        if not messages:
            return jsonify({"error": "Messages empty"}), 400
    except Exception:
        return jsonify({"error": "Invalid JSON"}), 400

    # सिस्टम प्रॉम्प्ट सेट करना
    system_prompt = os.getenv(mode.upper())
    if not system_prompt:
        return jsonify({"error": f"Mode {mode} missing in Secrets."}), 500

    messages.insert(0, {"role": "system", "content": system_prompt})

    # स्ट्रीमिंग जेनरेटर फंक्शन
    def generate_stream():
        try:
            # आपके द्वारा दिए गए OpenAI SDK कोड का सीधा उपयोग
            completion = client.chat.completions.create(
                model=MODEL_ID,
                messages=messages,
                temperature=0.7,
                top_p=0.8,
                max_tokens=1024,
                stream=True
            )
            
            for chunk in completion:
                if chunk.choices and chunk.choices[0].delta.content is not None:
                    # Frontend (index.html) को बिल्कुल वही JSON स्ट्रक्चर चाहिए जो NVIDIA देता था
                    # chunk.model_dump_json() इसे एकदम सही फॉर्मेट में बदल देता है
                    chunk_data = chunk.model_dump_json()
                    yield f"data: {chunk_data}\n\n"
                    
        except Exception as e:
            logger.error(f"OpenAI API Error: {str(e)}")
            yield f"data: {{\"error\": \"AI Provider Error: {str(e)}\"}}\n\n"

    return Response(stream_with_context(generate_stream()), mimetype='text/event-stream')

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