import os
import time
import json
import logging
import requests
import traceback
from typing import Dict, Any, List, Generator
from flask import Flask, request, Response, stream_with_context, jsonify

# --- सेटअप: लॉगिंग और कॉन्फिग ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - ARJUN-CORE-V3 - %(levelname)s - %(message)s')
logger = logging.getLogger("ArjunCoreV3")

app = Flask(__name__)

# कॉन्फ़िगरेशन मैनेजमेंट (बिना किसी हार्डकोडिंग के)
class Config:
    API_KEY = os.environ.get("YOUR_VEDIKA_API_KEY")
    ENDPOINT = os.environ.get("BASE_URL")
    MODEL_ID = os.environ.get("MODEL_ID")

    @classmethod
    def is_ready(cls):
        return all([cls.API_KEY, cls.ENDPOINT, cls.MODEL_ID])

# --- पेलोड बिल्डर: मॉडल की ज़रूरत के अनुसार डेटा बनाना ---
class PayloadBuilder:
    @staticmethod
    def construct(user_msg: str, history: List[Dict], system_prompt: str, attachments: List[Dict]):
        messages = [{"role": "system", "content": system_prompt or "You are a helpful assistant."}]
        
        # हिस्ट्री जोड़ना
        for h in history:
            messages.append({"role": h.get("role", "user"), "content": h.get("content", "")})
            
        # मल्टीमॉडल या टेक्स्ट पेलोड का चयन
        if attachments:
            content = [{"type": "text", "text": user_msg}]
            for att in attachments:
                if att.get("type") == "image":
                    content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{att.get('data')}"}})
            messages.append({"role": "user", "content": content})
        else:
            messages.append({"role": "user", "content": user_msg})
        return messages

# --- स्ट्रीम हैंडलर: LLM से रिस्पॉन्स लाना ---
def stream_llm_response(payload: Dict):
    headers = {"Authorization": f"Bearer {Config.API_KEY}", "Content-Type": "application/json"}
    try:
        with requests.post(Config.ENDPOINT, headers=headers, json=payload, stream=True, timeout=60) as r:
            if r.status_code != 200:
                yield f"data: {json.dumps({'error': 'API Error', 'code': r.status_code})}\n\n"
                return
            for line in r.iter_lines():
                if line:
                    yield line.decode('utf-8') + "\n\n"
    except Exception as e:
        logger.error(f"Stream Error: {e}")
        yield f"data: {json.dumps({'error': 'Connection Lost'})}\n\n"

# --- API रूट्स ---
@app.route('/api/chat', methods=['POST'])
def chat():
    if not Config.is_ready():
        return jsonify({"error": "Server not configured"}), 500
    
    data = request.json
    messages = PayloadBuilder.construct(
        data.get("message", ""), 
        data.get("history", []), 
        data.get("system_prompt", ""),
        data.get("attachments", [])
    )
    
    payload = {
        "model": Config.MODEL_ID,
        "messages": messages,
        "stream": True,
        "temperature": data.get("temperature", 0.7)
    }
    
    return Response(stream_with_context(stream_llm_response(payload)), mimetype='text/event-stream')

@app.route('/health', methods=['GET'])
def health():
    return jsonify({"status": "Arjun Core V3 Online", "timestamp": time.time()})

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 7860))
    logger.info(f"Arjun Core V3 starting on port {port}...")
    app.run(host='0.0.0.0', port=port, threaded=True)