# core/api_clients.py
import os
import time
import random
import requests
import shutil
import subprocess
import uuid
import base64
from threading import Lock
from gradio_client import Client

class AwaazAPI:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://awaaz-j36o.onrender.com/api"
        self.headers = {
            "X-API-Key": self.api_key,
            "Content-Type": "application/json"
        }

    def enhance_script(self, text):
        print("-> 🎭 Awaaz API से स्क्रिप्ट में इमोशन टैग्स जोड़े जा रहे हैं...")
        try:
            response = requests.post(f"{self.base_url}/proxy-enhance", json={"text": text}, headers=self.headers, timeout=60)
            response.raise_for_status()
            enhanced = response.json().get("enhanced_text", text)
            print("-> ✅ Script Enhanced!")
            return enhanced
        except Exception as e:
            print(f"🚨 Awaaz Enhance Error (Skipping Enhancement): {e}")
            return text

    def generate_audio(self, text, output_path):
        print("-> 🎙️ Awaaz API (Custom AI TTS) से उच्च-गुणवत्ता ऑडियो बनाया जा रहा है...")
        try:
            response = requests.post(f"{self.base_url}/proxy-tts", json={"text": text}, headers=self.headers, timeout=300)
            if response.status_code == 403:
                raise Exception("Awaaz API Key is invalid! (403 Forbidden)")
            response.raise_for_status()
            with open(output_path, "wb") as f:
                f.write(response.content)
            print(f"-> ✅ Custom AI Audio Saved: {output_path}")
            return output_path
        except Exception as e:
            raise Exception(f"🚨 Awaaz TTS Generation Error: {e}")

class GroqAPI:
    def __init__(self, api_keys):
        self.api_keys = api_keys
        self.audio_url = "https://api.groq.com/openai/v1/audio/transcriptions"
        self.chat_url = "https://api.groq.com/openai/v1/chat/completions"
        self.audio_model = "whisper-large-v3"
        self.chat_model = "meta-llama/llama-4-scout-17b-16e-instruct"
        self._key_index = 0
        self._lock = Lock()

    def get_next_key(self):
        with self._lock:
            key = self.api_keys[self._key_index % len(self.api_keys)]
            self._key_index += 1
            return key

    def transcribe_audio(self, audio_path):
        if not self.api_keys:
            raise Exception("Groq API key not found.")
        for attempt in range(len(self.api_keys) * 2):
            api_key = self.get_next_key()
            data = {'model': self.audio_model, 'response_format': 'verbose_json', 'timestamp_granularities[]': 'word'}
            headers = {'Authorization': f'Bearer {api_key}'}
            try:
                with open(audio_path, 'rb') as audio_file:
                    mime_type = 'audio/wav' if audio_path.endswith('.wav') else 'audio/mpeg'
                    files = {'file': (os.path.basename(audio_path), audio_file, mime_type)}
                    print(f"-> Groq API को ट्रांसक्रिप्शन भेजा जा रहा है...")
                    response = requests.post(self.audio_url, headers=headers, data=data, files=files, timeout=120)
                    if response.status_code == 429:
                        time.sleep(2 ** (attempt % 3))
                        continue
                    response.raise_for_status()
                    return response.json().get('words', [])
            except Exception as e:
                time.sleep(1)
        raise Exception("Groq Transcription Error: All keys failed.")

    def translate_to_english(self, hindi_text):
        if not hindi_text or hindi_text == "[PAUSE]":
            return ""
        api_key = self.get_next_key()
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        system_prompt = "You are a highly accurate Hindi to English translator. Translate the given Hindi text to English. Output ONLY the English translation, without any quotes, explanations, or extra text."
        payload = {
            "model": self.chat_model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": hindi_text}
            ],
            "temperature": 0.3,
            "max_tokens": 100
        }
        try:
            response = requests.post(self.chat_url, headers=headers, json=payload, timeout=15)
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content'].strip().strip('"').strip("'")
        except Exception as e:
            print(f"🚨 Groq Translation Error: {e}")
            return hindi_text

    def generate_visual_prompt(self, script_line, media_type):
        api_key = self.get_next_key()
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        if media_type == "png":
            rules = "Write a LITERAL prompt for an ISOLATED object or character against a SOLID WHITE background (for easy removal). Focus solely on the subject details, textures, and lighting. Describe it as a highly detailed cutout. NO scene background, NO context, NO metaphors."
        elif media_type == "video":
            rules = "Write a highly detailed, LITERAL text-to-video prompt IN ENGLISH ONLY. Describe motion, camera angle, lighting, and subject exactly. NO metaphors. NO split screens, NO text."
        else:
            rules = "Write a highly detailed, LITERAL oil-painting style prompt IN ENGLISH ONLY. Focus on dramatic lighting, epic scale, and rich colors. NO metaphors. NO frames, NO text."

        system_prompt = f"You are a strict Literal Visual Prompt Engineer. {rules} Output ONLY the English prompt string, nothing else."
        payload = {
            "model": self.chat_model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Scene context: {script_line}"}
            ],
            "temperature": 0.7,
            "max_tokens": 100
        }
        try:
            response = requests.post(self.chat_url, headers=headers, json=payload, timeout=30)
            response.raise_for_status()
            prompt_result = response.json()['choices'][0]['message']['content'].strip()
            return prompt_result.strip('"').strip("'")
        except Exception as e:
            print(f"🚨 Groq Prompt Engineer Error: {e}")
            return script_line

class StockClipAPI:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://scfai-f.vercel.app"
        self.headers = {
            "X-API-Key": self.api_key,
            "Content-Type": "application/json"
        }
        self.ui_logger = None

    def set_logger(self, logger_func):
        self.ui_logger = logger_func

    # 💡 डिफ़ॉल्ट प्रोग्रेस 35 सेट कर दिया है, ताकि डेटाबेस कभी खाली (NULL) न जाए
    def log(self, message, progress=35):
        if self.ui_logger:
            self.ui_logger(message, progress)
        else:
            print(message)

    def search_and_download_batch(self, query, download_dir, chunk_idx, scene_index, orientation="vertical", quality="1080p+", top_n=3):
        """
        सर्च इनिशिएट करता है और टॉप N क्लिप्स को डाउनलोड करके उनके पाथ्स की लिस्ट देता है।
        अगर 90 सेकंड तक रिजल्ट नहीं मिलता, तो फॉलबैक के लिए खाली लिस्ट रिटर्न करता है।
        """
        self.log(f"-> 🌐 StockClip AI: '{query}' के लिए हाई-क्वालिटी क्लिप्स खोजी जा रही हैं...", 32)
        
        # API डॉक्यूमेंटेशन के हिसाब से ओरिएंटेशन सेट करना
        ori = "portrait" if orientation == "vertical" else "landscape"
        payload = {
            "query": query,
            "orientation": ori,
            "quality": quality
        }
        
        try:
            # 1. 🚀 Task Initiate करना
            init_resp = requests.post(f"{self.base_url}/api/search", json=payload, headers=self.headers, timeout=60)
            init_resp.raise_for_status()
            task_id = init_resp.json().get("task_id")
            
            if not task_id:
                self.log("🚨 StockClip AI Error: Task ID नहीं मिला!", 32)
                return []
                
            self.log(f"-> ⏳ StockClip Task [{task_id[:6]}] शुरू! क्लिप्स ढूँढी जा रही हैं...", 33)
            
            # लॉग स्पैम रोकने के लिए ट्रैकर
            last_logged_status = None 
            
            # 💡 Hard Timeout सेटअप (90 सेकंड = 1.5 मिनट)
            start_time = time.time()
            max_polling_time = 160  
            
            # 2. 🔄 Task Status Polling
            while True:
                # 🛑 Timeout Check: अगर 90 सेकंड से ज़्यादा हो गया, तो Fallback ट्रिगर करो
                elapsed_time = time.time() - start_time
                if elapsed_time > max_polling_time:
                    self.log(f"🚨 StockClip AI Timeout Error: 90 सेकंड तक कोई जवाब नहीं मिला! Fallback ट्रिगर कर रहे हैं...", 35)
                    return []  # खाली लिस्ट भेजने से तुम्हारा Meta Shield अपने आप ट्रिगर हो जाएगा
                    
                time.sleep(5.5) # Polling Interval
                
                try:
                    status_resp = requests.get(f"{self.base_url}/api/status/{task_id}", headers=self.headers, timeout=60)
                    
                    # 504 Gateway Timeout (Cold Boot) Handling
                    if status_resp.status_code == 504:
                        if last_logged_status != "504":
                            self.log("-> 🚦 AI Backend Cold Boot (504). सर्वर वार्म-अप हो रहा है...", 33)
                            last_logged_status = "504"
                        continue
                        
                    # 401 Unauthorized
                    if status_resp.status_code == 401:
                        self.log("🚨 StockClip AI Auth Error (401)! API Key चेक करें।", 33)
                        return []
                        
                    # 404 Not Found (Task Syncing)
                    if status_resp.status_code == 404:
                        if last_logged_status != "404":
                            self.log("-> 🔄 Task सर्वर पर प्रोसेस हो रहा है... कृपया प्रतीक्षा करें...", 33)
                            last_logged_status = "404"
                        continue

                    # सफलता या अन्य सर्वर रिस्पांस
                    status_resp.raise_for_status()
                    data = status_resp.json()
                    task_info = data.get("task", {})
                    status = task_info.get("status")
                    
                    if status == "completed":
                        clips_data = task_info.get("data", [])
                        if not clips_data:
                            self.log(f"-> ⚠️ StockClip AI: '{query}' के लिए कोई क्लिप नहीं मिली।", 35)
                            return []
                            
                        self.log(f"-> ✅ StockClip AI ने {len(clips_data)} क्लिप्स ढूंढीं! टॉप {top_n} डाउनलोड की जा रही हैं...", 35)
                        
                        downloaded_paths = []
                        # सिर्फ टॉप N क्लिप्स डाउनलोड करें (Gemini के चुनने के लिए)
                        for i, clip in enumerate(clips_data[:top_n]):
                            dl_url = clip.get("download_url")
                            if not dl_url: continue
                            
                            clip_path = os.path.join(download_dir, f"chunk_{chunk_idx}_scene_{scene_index+1}_stockclip_{i}.mp4")
                            dl_resp = requests.get(dl_url, stream=True, timeout=120)
                            dl_resp.raise_for_status()
                            
                            with open(clip_path, 'wb') as f:
                                for chunk in dl_resp.iter_content(chunk_size=8192):
                                    f.write(chunk)
                                    
                            downloaded_paths.append(clip_path)
                            
                        return downloaded_paths
                        
                    elif status in ["failed", "error"]:
                        self.log("🚨 StockClip AI Task Failed!", 35)
                        return []
                        
                    else:
                        # queued, processing, fetching, ranking, vision
                        if last_logged_status != status:
                            self.log(f"-> 🔄 StockClip AI Status: {str(status).upper()}...", 34)
                            last_logged_status = status
                            
                except requests.exceptions.RequestException as e:
                    if last_logged_status != "network_error":
                        self.log(f"⚠️ StockClip Network Wait: सर्वर से जुड़ रहे हैं...", 34)
                        last_logged_status = "network_error"
                        
        except Exception as e:
            self.log(f"🚨 StockClip AI Generate Error: {str(e)}", 35)
            return []

class HuggingFacePNGAPI:
    def __init__(self, space_name, master_key):
        self.space_name = space_name
        self.master_key = master_key
        self.client = None
        self._lock = Lock()

    def generate_and_download(self, prompt, output_path):
        print(f"-> Gradio Client के ज़रिए HF Space '{self.space_name}' से AI PNG बनाया जा रहा है: '{prompt}'")
        try:
            with self._lock:
                if self.client is None:
                    print(f"-> 🔌 PNG API पहली बार कनेक्ट हो रहा है ({self.space_name})...")
                    self.client = Client(self.space_name)
                    
            result_path = self.client.predict(
                prompt=prompt, 
                secret_key=self.master_key, 
                api_name="/generate_png"
            )
            shutil.copy(result_path, output_path)
            print(f"-> ✅ AI PNG सफलतापूर्वक सहेजा गया: {output_path}")
            return output_path
        except Exception as e:
            print(f"🚨 AI PNG जनरेशन में त्रुटि (Gradio Client): {e}")
            return None

# ============================================================================== #
# 3. Custom AI Video/Image Generators (Meta Spark Studio & Sparkling API)
# ============================================================================== #

class CustomImageGenAPI:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://imgen-f.vercel.app"
        self.headers = {
            "X-API-Key": self.api_key,
            "Content-Type": "application/json"
        }

    def _save_base64_image(self, b64_string, file_path):
        """Base64 स्ट्रिंग को डिकोड करके फाइल के रूप में सेव करने का हेल्पर"""
        if "," in b64_string:
            b64_string = b64_string.split(",")[1]
        image_data = base64.b64decode(b64_string)
        with open(file_path, "wb") as f:
            f.write(image_data)

    def generate_and_download(self, prompt, download_path, orientation="vertical"):
        print(f"-> 🎨 Sparkling API से AI Image बनाई जा रही है: '{prompt[:50]}...'")
        
        # 💡 API Documentation के हिसाब से रेश्यो सेट करना
        ratio_str = "9:16" if orientation == "vertical" else "16:9"
        
        payload = {
            "prompt": prompt,
            "user_negative": "text, watermark, ugly, deformed, blurry, bad anatomy",
            "style_name": "Oil Painting",  # "Cinematic", "Anime", "Cosmic" भी इस्तेमाल कर सकते हो
            "ratio": ratio_str
        }
        
        try:
            # 1. 🚀 Task Initiate करना (Headers के साथ)
            response = requests.post(f"{self.base_url}/api/generate", json=payload, headers=self.headers, timeout=160)
            response.raise_for_status()
            data = response.json()
            
            # ⚡ Magic Cache Hit (अगर इमेज पहले से मौजूद है)
            if data.get("cached") and data.get("image"):
                print("-> ⚡ Magic Cache Hit! Sparkling API ने तुरंत इमेज लौटा दी।")
                self._save_base64_image(data["image"], download_path)
                return download_path
                
            task_id = data.get("task_id")
            if not task_id:
                print("🚨 Sparkling Image API Error: Task ID नहीं मिला!")
                return None
                
            print(f"-> ⏳ Task [{task_id[:6]}] कतार में है। Status Polling शुरू...")
            
            # 2. 🔄 Task Status Polling (Headers के साथ)
            while True:
                status_res = requests.get(f"{self.base_url}/api/status/{task_id}", headers=self.headers, timeout=160)
                
                if status_res.status_code == 404:
                    print("🚨 Sparkling API Task Error: Task 404 Not Found या Expire हो गया!")
                    return None
                    
                status_res.raise_for_status()
                res_data = status_res.json()
                status = res_data.get("status")
                
                if status == "completed":
                    print("-> ✅ Sparkling API Image सफलतापूर्वक जनरेट हो गई!")
                    self._save_base64_image(res_data["image"], download_path)
                    return download_path
                    
                elif status in ["failed", "error"]:
                    print(f"🚨 Sparkling API Task Failed: {res_data.get('error', 'Unknown Error')}")
                    return None
                    
                else:
                    # 'queued' या 'processing' स्टेटस में 2.5 सेकंड रुकना
                    time.sleep(2.5)
                    
        except Exception as e:
            print(f"🚨 Custom Image Gen API Error: {e}")
            return None

class MetaSparkStudioAPI:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://meta-genai.vercel.app".rstrip("/")
        self.headers = {
            "x-api-key": self.api_key,
            "Content-Type": "application/json"
        }
        self.ui_logger = None

    def set_logger(self, logger_func):
        self.ui_logger = logger_func

    def log(self, message):
        if self.ui_logger:
            self.ui_logger(message, 60)
        else:
            print(message)

    def _process_and_download_urls(self, urls, download_path, is_image_task):
        if not urls:
            self.log("🚨 MSS API Error: URL लिस्ट खाली है!")
            return None

        if is_image_task:
            self.log(f"-> 📥 MSS ने {len(urls)} इमेजेस दी हैं। पहली वाली को चुना जा रहा है...")
            dl = requests.get(urls[0], stream=True, timeout=None)
            dl.raise_for_status()
            with open(download_path, 'wb') as f:
                for chunk in dl.iter_content(chunk_size=8192):
                    f.write(chunk)
            self.log("-> ✅ 🛡️ Ultimate Shield: 1 शानदार इमेज सहेजी गई!")
            return download_path
            
        elif len(urls) >= 2:
            self.log(f"-> 📥 MSS ने {len(urls)} क्लिप्स दी हैं। पहली दो (5s+5s) को फेविकोल से चिपकाया जा रहा है...")
            base_dir = os.path.dirname(download_path)
            
            uid = uuid.uuid4().hex[:6]
            vid1_path = os.path.join(base_dir, f"mss_part1_{uid}.mp4")
            vid2_path = os.path.join(base_dir, f"mss_part2_{uid}.mp4")
            
            dl1 = requests.get(urls[0], stream=True, timeout=None)
            dl1.raise_for_status()
            with open(vid1_path, 'wb') as f:
                for chunk in dl1.iter_content(chunk_size=8192): f.write(chunk)
                
            dl2 = requests.get(urls[1], stream=True, timeout=None)
            dl2.raise_for_status()
            with open(vid2_path, 'wb') as f:
                for chunk in dl2.iter_content(chunk_size=8192): f.write(chunk)
                
            concat_txt_path = os.path.join(base_dir, f"mss_concat_{uid}.txt")
            with open(concat_txt_path, 'w') as f:
                f.write(f"file '{os.path.abspath(vid1_path)}'\n")
                f.write(f"file '{os.path.abspath(vid2_path)}'\n")
                
            cmd = ['ffmpeg', '-y', '-f', 'concat', '-safe', '0', '-i', concat_txt_path, '-c', 'copy', download_path]
            subprocess.run(cmd, capture_output=True, check=True)
            
            try:
                os.remove(vid1_path)
                os.remove(vid2_path)
                os.remove(concat_txt_path)
            except: 
                pass
            
            self.log("-> ✅ 🛡️ Ultimate Shield: दो 5s क्लिप्स सफलता पूर्वक 10s में जुड़ गईं!")
            return download_path
            
        elif len(urls) == 1:
            dl = requests.get(urls[0], stream=True, timeout=None)
            dl.raise_for_status()
            with open(download_path, 'wb') as f:
                for chunk in dl.iter_content(chunk_size=8192): f.write(chunk)
            self.log("-> ✅ 🛡️ Ultimate Shield: केवल 1 क्लिप मिली, सहेजी गई!")
            return download_path

    def generate_and_download(self, prompt, download_path, orientation="vertical"):
        short_prompt = prompt[:30] + "..." if prompt else "..."
        self.log(f"-> 🛡️ MSS API (Ultimate Shield) को रिक्वेस्ट: '{short_prompt}'")
        
        is_image_task = "STATIC IMAGE" in prompt
        exact_ratio = "9:16" if orientation == "vertical" else "16:9"
        mss_prompt = f"Generate a video in {exact_ratio} ratio. {prompt}" if not is_image_task else f"Generate an image in {exact_ratio} ratio. {prompt}"
        
        payload = {
            "prompt": mss_prompt,
            "quality": "high",
            "duration": 10
        }
        
        try:
            init_resp = requests.post(f"{self.base_url}/api/generate", json=payload, headers=self.headers, timeout=60)
            init_resp.raise_for_status()
            init_data = init_resp.json()
            
            direct_urls = init_data.get("urls") or init_data.get("result", {}).get("urls")
            if direct_urls:
                self.log("-> ⚡ MSS Cache Hit! सीधे जनरेटेड फाइल्स मिल गईं।")
                return self._process_and_download_urls(direct_urls, download_path, is_image_task)
            
            task_id = init_data.get("task_id") or init_data.get("id")
            
            if not task_id:
                self.log(f"🚨 MSS API Error: Task ID या URLs दोनों नहीं मिले! Response: {init_data}")
                return None
                
            self.log(f"-> ⏳ MSS Task [{task_id[:6]}] आरंभ हुआ। Polling शुरू...")
            wait_time = 5
            
            while True:
                try:
                    status_resp = requests.get(f"{self.base_url}/api/status/{task_id}", headers=self.headers, timeout=60)
                    
                    if status_resp.status_code == 404:
                        self.log("🚨 MSS Task Error: Task 404 Not Found या Expire हो गया!")
                        return None
                        
                    status_resp.raise_for_status()
                    data = status_resp.json()
                    
                    status = data.get("status")
                    
                    if status == "completed" or "urls" in data or ("result" in data and "urls" in data["result"]):
                        urls = data.get("urls") or data.get("result", {}).get("urls", [])
                        return self._process_and_download_urls(urls, download_path, is_image_task)
                            
                    elif status in ["failed", "error"]:
                        self.log(f"🚨 MSS Task Failed: {data.get('error', 'Unknown Error')}")
                        return None
                    else:
                        progress = data.get("progress", 0)
                        self.log(f"-> 🔄 MSS Status: {str(status).upper()} | Progress: {progress}%")
                        time.sleep(wait_time)
                        wait_time = min(wait_time * 1.5, 30)
                        
                except Exception as poll_err:
                    self.log(f"⚠️ MSS Status Check Error: {poll_err}. Retrying in {wait_time}s...")
                    time.sleep(wait_time)
                    wait_time = min(wait_time * 1.5, 30)
                    
        except Exception as e:
            self.log(f"🚨 MSS Generate API Error: {str(e)}")
            return None