import os
import requests
from datetime import datetime, timedelta, timezone
from flask import Flask, request, Response, stream_with_context, render_template_string

app = Flask(__name__)

# 🔐 --- SECURE ENVIRONMENT VARIABLES ---
API_KEY = os.environ.get("NVIDIA_API_KEY") or os.environ.get("YOUR_VEDIKA_API_KEY") 
MODEL_ID = os.environ.get("MODEL_ID", "google/diffusiongemma-26b-a4b-it") # Default fallback

# NVIDIA's official invoke URL
INVOKE_URL = "https://integrate.api.nvidia.com/v1/chat/completions"

# 🔑 --- SERPAPI KEY (Provided by Divy Patel) ---
SERPAPI_KEY = "df7dc67448cc9fe63ee6bbc7c20c50662ac18bc1f4682cd480da97bde1970381"

# 🌐 --- SERPAPI GOOGLE SEARCH ENGINE (100% BULLETPROOF) --- 🌐
def web_search_scraper(query, num_results=5):
    """
    यह SerpApi का उपयोग करके सीधे Google Search से एकदम सटीक और ताज़ा (JSON) डेटा लाता है।
    इसे Hugging Face या Google कभी ब्लॉक नहीं कर सकता।
    """
    results = []
    try:
        params = {
            "engine": "google",
            "q": query,
            "api_key": SERPAPI_KEY,
            "num": num_results,
            "hl": "en",
            "gl": "in" # India Region for local context
        }
        
        response = requests.get("https://serpapi.com/search", params=params, timeout=10)
        data = response.json()
        
        # 'organic_results' से असली Google सर्च का डेटा निकालना
        if "organic_results" in data:
            for item in data["organic_results"]:
                title = item.get("title", "")
                link = item.get("link", "")
                snippet = item.get("snippet", "")
                
                if title and snippet:
                    results.append({
                        "title": title, 
                        "link": link, 
                        "snippet": snippet
                    })
                    
    except Exception as e:
        print(f"SerpApi Error: {e}")

    return results

# ----------------------------------------------------

@app.route('/')
def home():
    try:
        with open('index.html', 'r', encoding='utf-8') as f:
            return render_template_string(f.read())
    except Exception as e:
        return f"<h1>System Error</h1><p>index.html missing: {str(e)}</p>"

@app.route('/api/chat', methods=['POST'])
def chat():
    if not API_KEY or not INVOKE_URL or not MODEL_ID:
        return Response("Server Error: Secrets missing. Check NVIDIA API Key.", status=500)

    data = request.get_json() or {}
    user_message = data.get("message", "")
    attachments = data.get("attachments", []) 
    is_search = data.get("is_search", False)
    history = data.get("history", []) 
    max_tokens = data.get("max_tokens", 4096)
    temperature = data.get("temperature", 1.0) # NVIDIA Recommended

    # 🕒 --- REAL-TIME IST INJECTION ---
    ist_time = datetime.now(timezone.utc) + timedelta(hours=5, minutes=30)
    current_date = ist_time.strftime("%A, %d %B %Y, %I:%M %p IST")

    # 🧠 --- GOD MODE SYSTEM PROMPT ---
    system_prompt = f"""
    You are CODE VED, an advanced AI System engineered EXCLUSIVELY by DIVY PATEL.
    Current Live Date and Time: {current_date}.
    
    STRICT DIRECTIVES:
    1. NEVER invent, guess, or hallucinate product launches, dates, news, or facts.
    2. If you receive "LIVE WEB SEARCH RESULTS", you MUST base your answer ENTIRELY on that data.
    3. Do NOT say "Based on the provided search results". Just answer naturally and confidently, citing the sources/links if needed.
    """

    # 🚀 --- AUTO-SEARCH INJECTION (USING SERPAPI) ---
    if is_search:
        scraped_data = web_search_scraper(user_message)
        
        search_context = "\n\n--- [LIVE VERIFIED GOOGLE SEARCH DATA] ---\n"
        if scraped_data:
            for idx, res in enumerate(scraped_data):
                search_context += f"{idx+1}. TITLE: {res['title']}\nSNIPPET: {res['snippet']}\nURL: {res['link']}\n\n"
            
            search_context += "[SYSTEM COMMAND: Use the above live Google data to answer the user accurately. Synthesize the info naturally without claiming you read it from a prompt.]"
        else:
            search_context += "[SYSTEM ALERT: Live search did not return results. Rely on your existing knowledge, but DO NOT hallucinate recent news.]"
            
        user_message = f"USER QUERY: {user_message}\n\n{search_context}"

    messages = [{"role": "system", "content": system_prompt}]

    for msg in history:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        if content:
            messages.append({"role": role, "content": content})

    content_payload = []
    if user_message.strip():
        content_payload.append({"type": "text", "text": user_message})
        
    for att in attachments:
        att_type = att.get("type")
        b64_data = att.get("data")
        if att_type == "image":
            content_payload.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_data}"}})
        elif att_type in ["audio", "file"]:
            content_payload.append({"type": "input_audio", "input_audio": {"data": b64_data, "format": "wav"}})

    if not content_payload:
        content_payload.append({"type": "text", "text": "Hello"})

    messages.append({"role": "user", "content": content_payload})

    headers = {
        "Authorization": f"Bearer {API_KEY}", 
        "Accept": "text/event-stream"
    }

    # 🌟 --- NVIDIA FORMAT INTEGRATED ---
    payload = {
        "model": MODEL_ID,
        "messages": messages,
        "max_tokens": int(max_tokens),
        "temperature": float(temperature),
        "top_p": 0.95,
        "stream": True,
        "chat_template_kwargs": {"enable_thinking": True} # Thinking logic enabled
    }

    try:
        response = requests.post(INVOKE_URL, headers=headers, json=payload, stream=True)
        def generate():
            for line in response.iter_lines():
                if line:
                    decoded_line = line.decode("utf-8")
                    if decoded_line.startswith("data: "):
                        yield decoded_line + "\n\n"
        return Response(stream_with_context(generate()), mimetype='text/event-stream')
    except Exception as e:
        return Response(f"Internal Error: {str(e)}", status=500)

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