import os
import sys
import json
import time
import uuid
import logging
import asyncio
import tempfile
import threading
import traceback
import requests
from typing import Dict, Any, List, Optional, Generator

from flask import (
    Flask, 
    request, 
    Response, 
    stream_with_context, 
    render_template_string, 
    send_file, 
    jsonify,
    make_response
)
import edge_tts

# =========================================================================================
# SECTION 1: ADVANCED LOGGING CONFIGURATION
# =========================================================================================
# सिस्टम की हर गतिविधि को ट्रैक करने के लिए कस्टम लॉगर।
class CustomFormatter(logging.Formatter):
    """लॉगिंग आउटपुट को सुंदर और स्पष्ट बनाने के लिए कस्टम फॉर्मेटर।"""
    grey = "\x1b[38;20m"
    yellow = "\x1b[33;20m"
    red = "\x1b[31;20m"
    bold_red = "\x1b[31;1m"
    green = "\x1b[32;20m"
    reset = "\x1b[0m"
    format_str = "%(asctime)s - ARJUN-CORE - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"

    FORMATS = {
        logging.DEBUG: grey + format_str + reset,
        logging.INFO: green + format_str + reset,
        logging.WARNING: yellow + format_str + reset,
        logging.ERROR: red + format_str + reset,
        logging.CRITICAL: bold_red + format_str + reset
    }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(log_fmt, datefmt='%Y-%m-%d %H:%M:%S')
        return formatter.format(record)

logger = logging.getLogger("ArjunBackend")
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setFormatter(CustomFormatter())
logger.addHandler(ch)

# =========================================================================================
# SECTION 2: CUSTOM EXCEPTIONS
# =========================================================================================
# त्रुटियों को बारीकी से पकड़ने के लिए कस्टम एक्सेप्शन क्लासेस।
class ArjunBaseException(Exception):
    """अर्जुन सिस्टम के लिए बेस एक्सेप्शन क्लास।"""
    def __init__(self, message: str, status_code: int = 500):
        super().__init__(message)
        self.message = message
        self.status_code = status_code

class ConfigurationError(ArjunBaseException):
    """जब पर्यावरण चर (Environment Variables) गायब हों।"""
    pass

class LLMAPIConnectionError(ArjunBaseException):
    """जब भाषा मॉडल API से संपर्क टूट जाए।"""
    pass

class TTSGenerationError(ArjunBaseException):
    """जब प्रभात न्यूरल आवाज़ जनरेट करने में विफल हो।"""
    pass

# =========================================================================================
# SECTION 3: SYSTEM CONFIGURATION MANAGER
# =========================================================================================
class ConfigManager:
    """
    यह क्लास सिस्टम की सभी सेटिंग्स और सीक्रेट्स (Secrets) को मैनेज करती है।
    यह सुनिश्चित करती है कि कोई भी डेटा हार्डकोडेड (Hardcoded) न हो।
    """
    def __init__(self):
        logger.info("Initializing Configuration Manager...")
        self.api_key = os.environ.get("YOUR_VEDIKA_API_KEY")
        self.base_url = os.environ.get("BASE_URL")
        self.model_id = os.environ.get("MODEL_ID")
        
        self.invoke_url = self._format_base_url(self.base_url)
        self.validate_config()

    def _format_base_url(self, url: Optional[str]) -> str:
        """API URL को सही प्रारूप में बदलता है।"""
        if not url:
            return ""
        url = url.strip()
        if not url.endswith("/chat/completions"):
            return f"{url.rstrip('/')}/chat/completions"
        return url

    def validate_config(self):
        """जाँच करता है कि सभी ज़रूरी सीक्रेट्स मौजूद हैं या नहीं।"""
        missing = []
        if not self.api_key: missing.append("YOUR_VEDIKA_API_KEY")
        if not self.base_url: missing.append("BASE_URL")
        if not self.model_id: missing.append("MODEL_ID")
        
        if missing:
            error_msg = f"Critical Configuration Missing: {', '.join(missing)}"
            logger.error(error_msg)
            # हम सर्वर क्रैश नहीं करेंगे, बल्कि रिक्वेस्ट आने पर एरर देंगे।
        else:
            logger.info("All system configurations validated successfully.")

# ग्लोबल कॉन्फिग ऑब्जेक्ट
config = ConfigManager()

# =========================================================================================
# SECTION 4: TEMPORARY FILE MANAGER (GARBAGE COLLECTION)
# =========================================================================================
class TempFileManager:
    """
    यह क्लास सुनिश्चित करती है कि जनरेट की गई ऑडियो (MP3) फाइल्स 
    सिस्टम की मेमोरी न भरें। यह उन्हें सुरक्षित रूप से डिलीट करती है।
    """
    def __init__(self):
        self.temp_dir = tempfile.gettempdir()
        self.active_files = set()
        logger.info(f"Temporary File Manager initialized at {self.temp_dir}")

    def create_temp_audio_path(self) -> str:
        """एक यूनीक सुरक्षित फाइल पाथ जनरेट करता है।"""
        unique_id = uuid.uuid4().hex
        file_path = os.path.join(self.temp_dir, f"arjun_tts_{unique_id}.mp3")
        self.active_files.add(file_path)
        return file_path

    def schedule_cleanup(self, file_path: str, delay_seconds: int = 60):
        """कुछ समय बाद फाइल को डिलीट करने के लिए बैकग्राउंड थ्रेड चलाता है।"""
        def _cleanup():
            time.sleep(delay_seconds)
            try:
                if os.path.exists(file_path):
                    os.remove(file_path)
                    logger.debug(f"Garbage Collection: Deleted {file_path}")
                if file_path in self.active_files:
                    self.active_files.remove(file_path)
            except Exception as e:
                logger.error(f"Cleanup failed for {file_path}: {str(e)}")

        cleanup_thread = threading.Thread(target=_cleanup, daemon=True)
        cleanup_thread.start()

file_manager = TempFileManager()

# =========================================================================================
# SECTION 5: ADVANCED TTS ENGINE (PRABHAT NEURAL HANDLER)
# =========================================================================================
class TTSEngine:
    """
    टेक्स्ट-टू-स्पीच (TTS) को हैंडल करने के लिए विशेष क्लास।
    यह Flask के सिंक्रोनस थ्रेड्स और Asyncio के बीच के टकराव को दूर करती है।
    """
    def __init__(self):
        # प्रभात न्यूरल (India Male) की सेटिंग्स
        # Pitch: -2% (प्राकृतिक भारीपन), Rate: +0% (सामान्य गति)
        self.voice = "en-IN-PrabhatNeural"
        self.pitch = "-2%"
        self.rate = "+0%"
        logger.info(f"TTS Engine initialized with Voice: {self.voice}")

    def generate_audio_sync(self, text: str, output_path: str) -> bool:
        """
        यह फंक्शन asyncio लूप को एक सुरक्षित आइसोलेटेड (Isolated) तरीके से चलाता है
        ताकि Flask सर्वर कभी हैंग न हो। यह यूज़र की सबसे बड़ी समस्या का समाधान है।
        """
        # टेक्स्ट की सफाई (Markdown हटाना)
        clean_text = self._clean_text_for_speech(text)
        if not clean_text:
            logger.warning("TTS Engine received empty text after cleaning.")
            return False

        logger.info(f"Starting TTS Generation for text length: {len(clean_text)}")

        # Async फंक्शन जो असल में edge-tts को कॉल करेगा
        async def _async_generate():
            try:
                communicate = edge_tts.Communicate(
                    text=clean_text, 
                    voice=self.voice, 
                    pitch=self.pitch, 
                    rate=self.rate
                )
                await communicate.save(output_path)
                return True
            except Exception as e:
                logger.error(f"Edge-TTS Core Error: {str(e)}")
                traceback.print_exc()
                return False

        # नया इवेंट लूप बनाकर चलाना (Thread-safe)
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            success = loop.run_until_complete(_async_generate())
            return success
        except Exception as e:
            logger.error(f"Event Loop Error during TTS: {str(e)}")
            return False
        finally:
            loop.close()

    def _clean_text_for_speech(self, text: str) -> str:
        """AI के आउटपुट से अनचाहे चिह्न हटाता है ताकि आवाज़ न फटे।"""
        import re
        # Markdown इमेजेज हटाना
        text = re.sub(r'!\[.*?\]\(.*?\)', '', text)
        # Markdown लिंक्स हटाना
        text = re.sub(r'\[.*?\]\(.*?\)', '', text)
        # <think> टैग्स हटाना
        text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
        # बोल्ड/इटैलिक चिह्न हटाना
        text = text.replace('*', '').replace('_', '').replace('`', '')
        return text.strip()

tts_engine = TTSEngine()

# =========================================================================================
# SECTION 6: LLM API STREAMING HANDLER
# =========================================================================================
class LLMAPIHandler:
    """
    यह क्लास यूज़र के मैसेज और फाइल्स को प्रोसेस करके 
    भाषा मॉडल (LLM) को भेजती है और स्ट्रीमिंग रिस्पॉन्स को वापस लाती है।
    """
    def __init__(self, cfg: ConfigManager):
        self.cfg = cfg

    def construct_messages(self, user_msg: str, attachments: List[Dict], sys_prompt: str, history: List[Dict]) -> List[Dict]:
        """पेलोड (Payload) का निर्माण करता है।"""
        messages = []
        
        # 1. System Prompt
        if sys_prompt and sys_prompt.strip():
            messages.append({"role": "system", "content": sys_prompt.strip()})

        # 2. History (पुराने मैसेज)
        for h_msg in history:
            role = h_msg.get("role", "user")
            content = h_msg.get("content", "")
            if content:
                messages.append({"role": role, "content": content})

        # 3. Current User Message & Attachments
        content_payload = []
        if user_msg and user_msg.strip():
            content_payload.append({"type": "text", "text": user_message.strip()})
            
        for att in attachments:
            att_type = att.get("type")
            b64_data = att.get("data")
            
            if not att_type or not b64_data:
                continue
                
            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", "video"]:
                # वीडियो को भी ऑडियो/फाइल की तरह ट्रीट किया जाएगा अगर मॉडल सपोर्ट करता है
                content_payload.append({
                    "type": "input_audio", 
                    "input_audio": {"data": b64_data, "format": "wav"}
                })

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

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

    def stream_response(self, payload: Dict[str, Any]) -> Generator[str, None, None]:
        """API से डेटा स्ट्रीम करता है और सुरक्षित रूप से यील्ड (Yield) करता है।"""
        headers = {
            "Authorization": f"Bearer {self.cfg.api_key}",
            "Accept": "text/event-stream",
            "Content-Type": "application/json"
        }

        logger.info(f"Initiating streaming request to {self.cfg.invoke_url}...")
        
        try:
            with requests.post(
                self.cfg.invoke_url, 
                headers=headers, 
                json=payload, 
                stream=True,
                timeout=120 # 2 minute timeout for safety
            ) as response:
                
                # Check for HTTP errors
                if response.status_code != 200:
                    error_text = response.text
                    logger.error(f"API HTTP Error {response.status_code}: {error_text}")
                    yield f"data: {json.dumps({'choices': [{'delta': {'content': f'**API Error:** {response.status_code}'}}]})}\n\n"
                    yield "data: [DONE]\n\n"
                    return

                for line in response.iter_lines():
                    if line:
                        decoded_line = line.decode("utf-8")
                        if decoded_line.startswith("data: "):
                            yield decoded_line + "\n\n"
                            
        except requests.exceptions.Timeout:
            logger.error("API Request Timed Out.")
            yield f"data: {json.dumps({'choices': [{'delta': {'content': '**Error:** Request timed out.'}}]})}\n\n"
        except requests.exceptions.RequestException as e:
            logger.error(f"API Network Error: {str(e)}")
            yield f"data: {json.dumps({'choices': [{'delta': {'content': '**Error:** Network connection failed.'}}]})}\n\n"
        except Exception as e:
            logger.error(f"Unknown Streaming Error: {str(e)}")
            traceback.print_exc()
            yield f"data: {json.dumps({'choices': [{'delta': {'content': '**Error:** Internal server error during stream.'}}]})}\n\n"

llm_handler = LLMAPIHandler(config)

# =========================================================================================
# SECTION 7: FLASK APPLICATION SETUP
# =========================================================================================
app = Flask(__name__)

# Request Logging Middleware
@app.before_request
def log_request_info():
    """हर रिक्वेस्ट की जानकारी लॉग करता है।"""
    logger.debug(f"Incoming Request: {request.method} {request.path}")

@app.after_request
def add_cors_headers(response):
    """CORS और सिक्योरिटी हेडर्स जोड़ता है।"""
    response.headers["Access-Control-Allow-Origin"] = "*"
    response.headers["Access-Control-Allow-Methods"] = "POST, GET, OPTIONS"
    response.headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization"
    return response

# =========================================================================================
# SECTION 8: API ROUTES & ENDPOINTS
# =========================================================================================

@app.route('/', methods=['GET'])
def serve_frontend():
    """मुख्य HTML फाइल (index.html) सर्व करता है।"""
    logger.info("Serving index.html to client.")
    try:
        # सुनिश्चित करें कि index.html उसी डायरेक्टरी में है
        file_path = os.path.join(os.path.dirname(__file__), 'index.html')
        with open(file_path, 'r', encoding='utf-8') as f:
            html_content = f.read()
        return render_template_string(html_content)
    except FileNotFoundError:
        logger.error("index.html file not found on server.")
        return make_response("<h1>Error: index.html missing!</h1><p>Please place index.html in the same directory as app.py</p>", 404)
    except Exception as e:
        logger.error(f"Error reading index.html: {str(e)}")
        return make_response(f"<h1>Internal Error</h1><p>{str(e)}</p>", 500)

@app.route('/api/chat', methods=['POST', 'OPTIONS'])
def chat_endpoint():
    """LLM चैट प्रोसेसिंग के लिए मुख्य एंडपॉइंट।"""
    if request.method == 'OPTIONS':
        return Response(status=200)

    # 1. Validate Config
    if not config.api_key or not config.invoke_url or not config.model_id:
        logger.critical("Chat request denied. Missing Configuration Secrets.")
        return jsonify({"error": "Server is improperly configured. Check Secrets."}), 500

    # 2. Parse Request
    try:
        data = request.get_json() or {}
    except Exception as e:
        logger.error(f"Invalid JSON payload: {str(e)}")
        return jsonify({"error": "Invalid JSON format."}), 400

    user_message = data.get("message", "")
    attachments = data.get("attachments", []) 
    system_prompt = data.get("system_prompt", "")
    history = data.get("history", []) 
    max_tokens = data.get("max_tokens", 4096)
    temperature = data.get("temperature", 0.7)

    logger.info(f"Received chat request. Msg length: {len(user_message)}, Attachments: {len(attachments)}")

    # 3. Construct Payload
    try:
        messages = llm_handler.construct_messages(user_message, attachments, system_prompt, history)
    except Exception as e:
        logger.error(f"Payload construction failed: {str(e)}")
        return jsonify({"error": "Failed to construct message payload."}), 500

    api_payload = {
        "model": config.model_id,
        "messages": messages,
        "max_tokens": int(max_tokens),
        "temperature": float(temperature),
        "top_p": 0.70,
        "stream": True
    }

    # 4. Stream Response
    return Response(
        stream_with_context(llm_handler.stream_response(api_payload)), 
        mimetype='text/event-stream'
    )


@app.route('/api/tts', methods=['POST', 'OPTIONS'])
def tts_endpoint():
    """
    टेक्स्ट-टू-स्पीच (Prabhat Neural) जनरेशन एंडपॉइंट।
    यह ऑडियो फाइल बनाता है, उसे सर्व करता है, और फिर डिलीट करने का शेड्यूल बनाता है।
    """
    if request.method == 'OPTIONS':
        return Response(status=200)

    try:
        data = request.get_json() or {}
        text = data.get("text", "").strip()
    except Exception as e:
        logger.error(f"Invalid JSON in TTS request: {str(e)}")
        return jsonify({"error": "Invalid JSON payload"}), 400
    
    if not text:
        logger.warning("TTS Request received empty text.")
        return jsonify({"error": "No text provided"}), 400
        
    logger.info(f"Processing TTS request. Text length: {len(text)}")
    
    # एक सुरक्षित टेम्पररी फाइल पाथ प्राप्त करें
    output_audio_path = file_manager.create_temp_audio_path()
    
    # सिंक्रोनस थ्रेड के अंदर एसिंक्रोनस TTS को सुरक्षित रूप से चलाएं
    success = tts_engine.generate_audio_sync(text, output_audio_path)
    
    if not success or not os.path.exists(output_audio_path):
        logger.error("TTS Audio file was not created successfully.")
        return jsonify({"error": "TTS Engine failed to generate audio."}), 500
        
    logger.info(f"TTS Audio generated successfully at: {output_audio_path}")
    
    # फाइल को क्लाइंट को भेजने के बाद 60 सेकंड में डिलीट करने का शेड्यूल
    file_manager.schedule_cleanup(output_audio_path, delay_seconds=60)
    
    try:
        # फाइल को बाइट्स के रूप में भेजें (Streaming)
        return send_file(
            output_audio_path, 
            mimetype="audio/mpeg", 
            as_attachment=False,
            download_name="arjun_voice.mp3"
        )
    except Exception as e:
        logger.error(f"Failed to send TTS file: {str(e)}")
        return jsonify({"error": "Failed to transmit audio file."}), 500


@app.route('/health', methods=['GET'])
def health_check():
    """सर्वर की स्थिति जाँचने के लिए एंडपॉइंट।"""
    status = {
        "status": "healthy",
        "arjun_core": "online",
        "tts_engine": "initialized",
        "timestamp": time.time()
    }
    return jsonify(status), 200

# =========================================================================================
# SECTION 9: ERROR HANDLERS
# =========================================================================================
@app.errorhandler(404)
def not_found(error):
    logger.warning(f"404 Not Found: {request.path}")
    return jsonify({"error": "Resource not found"}), 404

@app.errorhandler(500)
def internal_error(error):
    logger.error(f"500 Internal Server Error: {request.path}")
    return jsonify({"error": "Internal Server Error"}), 500

# =========================================================================================
# MAIN EXECUTION BLOCK
# =========================================================================================
def print_startup_banner():
    """सर्वर शुरू होने पर एक सुंदर बैनर प्रिंट करता है।"""
    banner = """
    ========================================================
        █████╗ ██████╗      ██╗██╗   ██╗███╗   ██╗
       ██╔══██╗██╔══██╗     ██║██║   ██║████╗  ██║
       ███████║██████╔╝     ██║██║   ██║██╔██╗ ██║
       ██╔══██║██╔══██╗██   ██║██║   ██║██║╚██╗██║
       ██║  ██║██║  ██║╚█████╔╝╚██████╔╝██║ ╚████║
       ╚═╝  ╚═╝╚═╝  ╚═╝ ╚════╝  ╚═════╝ ╚═╝  ╚═══╝
               ADVANCED EDITION (V2.0)
    ========================================================
    System   : Arjun Intelligence Core
    TTS Mode : Microsoft Edge (Prabhat Neural)
    Status   : Online and Listening...
    ========================================================
    """
    print("\x1b[32;20m" + banner + "\x1b[0m")

if __name__ == '__main__':
    # बैनर दिखाएं
    print_startup_banner()
    
    # पोर्ट सेट करें
    port = int(os.environ.get("PORT", 7860))
    
    logger.info(f"Starting Arjun Backend Server on port {port}...")
    
    # प्रोडक्शन डिप्लॉयमेंट के लिए, Waitress या Gunicorn का इस्तेमाल बेहतर है,
    # लेकिन स्थानीय/डेवलपमेंट के लिए Flask का रन सर्वर (Threaded) पर्याप्त है।
    app.run(host='0.0.0.0', port=port, threaded=True)