import os
import time
import requests
import json
import base64
import threading
from PIL import Image

_model_lock = threading.Lock()

# Backend configuration via environment variable. Defaults to auto-detected or "mock"
try:
    import llama_cpp
    default_backend = "llama_cpp"
except ImportError:
    default_backend = "mock"
BACKEND = os.environ.get("BACKEND", default_backend).lower()

# Constants for Hugging Face Space model loading
MODEL_REPO = "bartowski/google_gemma-4-E2B-it-GGUF"
MODEL_FILE = "google_gemma-4-E2B-it-Q4_K_M.gguf"
LOCAL_MODEL_DIR = os.environ.get("MODEL_DIR", "./model")

_llama_model = None

def _download_gguf():
    """Download GGUF model from Hugging Face if not already present."""
    os.makedirs(LOCAL_MODEL_DIR, exist_ok=True)
    local_path = os.path.join(LOCAL_MODEL_DIR, MODEL_FILE)
    if os.path.exists(local_path):
        print(f"[llm.py] Model GGUF already exists at {local_path}")
        return local_path

    print(f"[llm.py] Downloading {MODEL_FILE} from HF repo {MODEL_REPO}...")
    try:
        from huggingface_hub import hf_hub_download
        downloaded_path = hf_hub_download(
            repo_id=MODEL_REPO,
            filename=MODEL_FILE,
            local_dir=LOCAL_MODEL_DIR,
            local_dir_use_symlinks=False
        )
        print(f"[llm.py] Model downloaded successfully to {downloaded_path}")
        return downloaded_path
    except Exception as e:
        print(f"[llm.py] Error downloading model from Hugging Face: {e}")
        return None

def init_llama_cpp():
    """Lazy initialization of llama_cpp model."""
    global _llama_model
    if _llama_model is not None:
        return _llama_model

    try:
        from llama_cpp import Llama
    except ImportError:
        print("[llm.py] Warning: llama-cpp-python is not installed. Falling back to mock backend.")
        return None

    model_path = _download_gguf()
    if not model_path or not os.path.exists(model_path):
        print("[llm.py] Error: Model file not found. Cannot load llama_cpp.")
        return None

    print(f"[llm.py] Loading model into memory: {model_path}")
    num_threads = 1 if os.environ.get("SPACE_ID") else 4
    try:
        _llama_model = Llama(
            model_path=model_path,
            n_ctx=2048,
            n_threads=num_threads,
            verbose=False
        )
        print("[llm.py] llama_cpp model loaded successfully!")
        return _llama_model
    except Exception as e:
        print(f"[llm.py] Error loading llama_cpp: {e}")
        return None

def downscale_image(image_path, max_dim=512):
    """
    Downscales the image if any dimension exceeds max_dim to save compute.
    Saves the downscaled image to a temporary file and returns its path.
    """
    if not image_path or not os.path.exists(image_path):
        return image_path
        
    try:
        img = Image.open(image_path)
        width, height = img.size
        if width > max_dim or height > max_dim:
            if width > height:
                new_width = max_dim
                new_height = int(height * (max_dim / width))
            else:
                new_height = max_dim
                new_width = int(width * (max_dim / height))
            
            img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
            dir_name = os.path.dirname(image_path)
            base_name = "downscaled_" + os.path.basename(image_path)
            temp_path = os.path.join(dir_name, base_name)
            img.save(temp_path, quality=85)
            print(f"[llm.py] Downscaled image from {width}x{height} to {new_width}x{new_height}")
            return temp_path
    except Exception as e:
        print(f"[llm.py] Error downscaling image: {e}")
    return image_path

def encode_image_to_base64(image_path):
    """Encodes image file to base64 string."""
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    except Exception as e:
        print(f"[llm.py] Error encoding image: {e}")
        return None

# --- Whisper.cpp ASR (Speech-to-Text) ---
_whisper_model = None

def _init_whisper():
    """Lazy initialization of whisper.cpp model for offline ASR."""
    global _whisper_model
    if _whisper_model is not None:
        return _whisper_model
    
    try:
        from pywhispercpp.model import Model as WhisperModel
        # Use multilingual 'tiny' model (~75MB) for Hindi/English support.
        # Auto-downloads GGML weights on first run, then fully offline.
        print("[llm.py] Loading whisper.cpp 'tiny' model for ASR...")
        _whisper_model = WhisperModel(
            'tiny',
            n_threads=2 if not os.environ.get("SPACE_ID") else 1
        )
        print("[llm.py] whisper.cpp ASR model loaded successfully!")
        return _whisper_model
    except ImportError:
        print("[llm.py] pywhispercpp not installed. ASR will use mock fallback.")
        return None
    except Exception as e:
        print(f"[llm.py] Error loading whisper.cpp ASR model: {e}")
        return None

def transcribe_audio(audio_path, prompt=""):
    """
    Transcribe audio file to text using whisper.cpp (offline, lightweight).
    Falls back to mock transcription if whisper.cpp is unavailable.
    """
    if not audio_path or not os.path.exists(audio_path):
        print("[llm.py] Audio file not found, using mock ASR.")
        return _mock_transcribe_audio(prompt)
    
    whisper = _init_whisper()
    if whisper is None:
        print("[llm.py] whisper.cpp unavailable, using mock ASR fallback.")
        return _mock_transcribe_audio(prompt)
    
    temp_wav_path = None
    try:
        # Transcode any audio format (MP3, FLAC, etc.) to 16kHz mono WAV using miniaudio
        try:
            import miniaudio
            import wave
            print(f"[llm.py] Decoding and resampling audio to 16kHz mono WAV using miniaudio...")
            sound = miniaudio.decode_file(audio_path, nchannels=1, sample_rate=16000)
            
            # Save to temp WAV file
            temp_wav_path = audio_path + ".temp_16k.wav"
            with wave.open(temp_wav_path, "wb") as wav_file:
                wav_file.setnchannels(1)
                wav_file.setsampwidth(2)  # 16-bit PCM
                wav_file.setframerate(16000)
                wav_file.writeframes(sound.samples)
            
            audio_path = temp_wav_path
            print(f"[llm.py] Resampled audio saved to: {audio_path}")
        except ImportError:
            print("[llm.py] miniaudio not installed. Passing audio file directly to whisper.cpp.")
        except Exception as e:
            print(f"[llm.py] miniaudio transcoding failed: {e}. Passing original file directly.")

        print(f"[llm.py] Transcribing audio: {audio_path}")
        segments = whisper.transcribe(audio_path)
        transcription = " ".join([seg.text.strip() for seg in segments]).strip()
        
        # Cleanup temp WAV file if created
        if temp_wav_path and os.path.exists(temp_wav_path):
            try: os.remove(temp_wav_path)
            except: pass
            
        if not transcription:
            print("[llm.py] Whisper returned empty transcription, using mock fallback.")
            return _mock_transcribe_audio(prompt)
        
        print(f"[llm.py] ASR Transcription: \"{transcription}\"")
        return transcription
    except Exception as e:
        # Cleanup temp WAV file on exception
        if temp_wav_path and os.path.exists(temp_wav_path):
            try: os.remove(temp_wav_path)
            except: pass
        print(f"[llm.py] Error during whisper.cpp transcription: {e}")
        return _mock_transcribe_audio(prompt)

def _mock_transcribe_audio(prompt=""):
    """Mock ASR fallback when whisper.cpp is not available."""
    prompt_lower = str(prompt).lower() if prompt else ""
    if "potato" in prompt_lower or "आलू" in prompt_lower:
        return "आलू में झुलसा रोग (blight) से कैसे बचाएं?"
    elif "wheat" in prompt_lower or "गेहूं" in prompt_lower:
        return "गेहूं में पीलापन आ रहा है, क्या करूँ?"
    else:
        return "फसल की सिंचाई करने का सही समय और तरीका क्या है?"

# Keep backward-compatible alias
mock_transcribe_audio = _mock_transcribe_audio

def generate_mock(prompt, system="", image_path=None, audio_path=None, history=None):
    """Simulate streaming for the mock backend."""
    import re
    response = ""
    
    # 0. Check if this is a structured data extraction query (from ledger)
    if "precise data extractor" in system.lower() or "json format" in system.lower():
        import datetime
        date_val = datetime.date.today().strftime("%Y-%m-%d")
        
        # Months pattern to avoid matching rupees or other words as dates
        months_pattern = r"(?:jan|feb|mar|apr|may|jun|jul|aug|sep|oct|nov|dec|january|february|march|april|june|july|august|september|october|november|december|जनवरी|फ़रवरी|मार्च|अप्रैल|मई|जून|जुलाई|अगस्त|सितंबर|अक्टूबर|नवंबर|दिसंबर)"
        date_match = re.search(r"(\d{1,2}\s+" + months_pattern + r"(?:\s+\d{4})?)", prompt, re.IGNORECASE)
        
        if date_match:
            date_val = date_match.group(1)
        else:
            # Check for standard date formats
            fmt_match = re.search(r"(\d{4}[-/]\d{1,2}[-/]\d{1,2}|\d{1,2}[-/]\d{1,2}[-/]\d{4})", prompt)
            if fmt_match:
                date_val = fmt_match.group(1)
            else:
                if any(w in prompt.lower() for w in ["yesterday", "kal", "कल"]):
                    date_val = (datetime.date.today() - datetime.timedelta(days=1)).strftime("%Y-%m-%d")
                elif any(w in prompt.lower() for w in ["today", "aaj", "आज"]):
                    date_val = datetime.date.today().strftime("%Y-%m-%d")
                
        price_val = 0
        numbers = re.findall(r"\b\d+\b", prompt)
        price_match = re.search(r"(\d+)\s*(?:rupees|rupee|rupay|rs|inr|me|में|रुपये|रुपए|रुपयों|रू|रु|\s|$)", prompt, re.IGNORECASE)
        if price_match:
            price_val = int(price_match.group(1))
        elif numbers:
            valid_nums = [int(n) for n in numbers if not (2020 <= int(n) <= 2040)]
            if valid_nums:
                price_val = valid_nums[-1]
                
        qty_val = "1 unit"
        qty_match = re.search(r"(\d+\s*(?:kilo|kg|kilos|kila|किलो|किग्रा|bags|bag|बोरी|बोरा|unit|units|लीटर|litre|l|liters))", prompt, re.IGNORECASE)
        if qty_match:
            qty_val = qty_match.group(1)
        elif numbers:
            valid_nums = [int(n) for n in numbers if not (2020 <= int(n) <= 2040) and int(n) != price_val]
            if valid_nums:
                qty_val = f"{valid_nums[0]} unit"
                
        item_val = "Crop"
        # Try to extract the item coming right after the quantity (e.g. "5 kg cotton", "1 bag urea")
        item_match = re.search(r"\b\d+\s*(?:kilo|kg|kilos|kila|किलो|किग्रा|bags|bag|बोरी|बोरा|unit|units|लीटर|litre|l|liters)\s+(?:of\s+)?([a-zA-Z\u0900-\u097F]+)", prompt, re.IGNORECASE)
        if item_match:
            item_val = item_match.group(1).strip()
        else:
            # Fallback to predefined lists
            items_list = ["rice", "wheat", "potato", "cotton", "urea", "fertilizer", "seed", "pesticide", "धान", "चावल", "गेहूं", "गेहूँ", "आलू", "कपास", "यूरिया", "खाद", "बीज", "कीटनाशक"]
            for it in items_list:
                if it in prompt.lower():
                    item_val = it
                    break
        
        type_val = "sale"
        if any(w in prompt.lower() for w in ["bought", "purchased", "खरीदा", "kharida", "khareeda", "kharid", "khareed", "buy", "purchase"]):
            type_val = "purchase"
        elif any(w in prompt.lower() for w in ["sold", "becha", "बेचा", "sale", "sell", "bech"]):
            type_val = "sale"
            
        data = {
            "date": date_val,
            "item": item_val,
            "qty": qty_val,
            "price": price_val,
            "type": type_val
        }
        
        yield json.dumps(data)
        return

    # 1. Handle Voice Audio ASR (whisper.cpp with mock fallback)
    if audio_path:
        transcription = transcribe_audio(audio_path, prompt)
        response += f"[🎙️ **आवाज का अनुवाद (ASR):** \"{transcription}\"]\n\n"
        prompt = transcription
        
    prompt_lower = prompt.lower()
    
    # 2. Handle Leaf Crop Image Mock Vision
    if image_path:
        response += (
            "📸 **छवि विश्लेषण (Image Analysis):**\n"
            "पत्ती की सतह पर भूरे-काले धब्बे और झुलसा हुआ भाग दिखाई दे रहा है। यह मुख्य रूप से **झुलसा रोग (Blight Disease)** का लक्षण प्रतीत होता है।\n\n"
        )
        if "potato" in prompt_lower or "आलू" in prompt_lower:
            response += (
                "**आलू पिछेती झुलसा रोग निवारण:**\n"
                "1. मैन्कोजेब (Mancozeb) का 2 ग्राम/लीटर पानी में मिलाकर छिड़काव करें।\n"
                "2. संक्रमित पौधों की पत्तियों को खेत से निकाल दें ताकि यह अन्य पौधों में न फैले।"
            )
        else:
            response += (
                "**गेहूं पीला रतुआ या धब्बा रोग निवारण:**\n"
                "1. प्रोपिकोनाजोल २५ ईसी का १ मिली प्रति लीटर पानी में छिड़काव करें।\n"
                "2. खेत में जल जमाव न होने दें और अतिरिक्त यूरिया के प्रयोग से बचें।"
            )
    else:
        # Standard text-based rule responses
        if "पीलापन" in prompt_lower or "yellow" in prompt_lower:
            response += (
                "गेहूं में पीलापन (Yellowing of wheat leaves) नाइट्रोजन की कमी या अधिक सिंचाई के कारण हो सकता है।\n\n"
                "**सलाह:**\n"
                "1. सिंचाई कम करें और मिट्टी में नमी की जांच करें।\n"
                "2. यूरिया (Nitrogen fertilizer) का छिड़काव करें (लगभग 20-25 किलोग्राम प्रति एकड़)।\n"
                "3. यदि समस्या बनी रहती है, तो 0.5% फेरस सल्फेट का छिड़काव करें।"
            )
        elif "झुलसा" in prompt_lower or "blight" in prompt_lower:
            response += (
                "आलू में झुलसा रोग (Blight disease) कवक (fungus) के कारण होता है। यह दो प्रकार का होता है: अगेती और पिछेती झुलसा।\n\n"
                "**सलाह:**\n"
                "1. मैन्коजेब (Mancozeb) या कॉपर ऑक्सीक्लोराइड का 2 ग्राम प्रति लीटर पानी में मिलाकर छिड़काव करें।\n"
                "2. संक्रमित पौधों की पत्तियों को खेत से निकाल दें ताकि यह अन्य पौधों में न फैले।\n"
                "3. रात में अधिक नमी होने पर विशेष ध्यान रखें।"
            )
        elif "सिंचाई" in prompt_lower or "irrigate" in prompt_lower or "water" in prompt_lower:
            response += (
                "गेहूं में सिंचाई की मुख्य 6 अवस्थाएं होती हैं:\n"
                "1. **ताज जड़ विकास (CRI)**: बुवाई के 21-25 दिन बाद (सबसे महत्वपूर्ण)।\n"
                "2. **कल्ले फूटते समय (Tillering)**: बुवाई के 40-45 दिन बाद।\n"
                "3. **गांठ बनते समय (Late Jointing)**: बुवाई के 60-65 दिन बाद।\n"
                "4. **फूल आने पर (Flowering)**: बुवाई के 80-85 दिन बाद।\n"
                "5. **दूधिया अवस्था (Milking)**: बुवाई के 100-105 दिन बाद।\n"
                "6. **दाना पकते समय (Dough)**: बुवाई के 115-120 दिन बाद।"
            )
        elif "ledger" in prompt_lower or "बहीखाता" in prompt_lower or "बेचा" in prompt_lower or "खरीदा" in prompt_lower or "sold" in prompt_lower or "bought" in prompt_lower:
            if "गेहूं" in prompt_lower or "wheat" in prompt_lower:
                item = "Wheat (गेहूं)"
                qty = "40 kg"
                price = "1200"
            elif "आलू" in prompt_lower or "potato" in prompt_lower:
                item = "Potato (आलू)"
                qty = "500 kg"
                price = "6000"
            else:
                item = "Crop (फसल)"
                qty = "100 kg"
                price = "2000"
                
            data = {
                "date": "Today",
                "item": item,
                "qty": qty,
                "price": price,
                "type": "sale" if ("बेचा" in prompt_lower or "sold" in prompt_lower) else "purchase"
            }
            response = f"JSON_OUTPUT: {json.dumps(data)}"
        elif any(kw in prompt_lower for kw in [
            "what were my last", "last message", "last question",
            "pichle message", "pichle sandesh", "pichle sawal",
            "pichla sawaal", "मेरे पिछले", "पिछले सवाल", "पिछला सवाल",
            "पिछले संदेश", "पिछला संदेश", "क्या पूछा था", "kya poocha tha",
            "mera pichla", "what did i ask", "what did i say"
        ]):
            if history:
                user_msgs = [msg["content"] for msg in history if msg.get("role") == "user"]
                if user_msgs:
                    numbered = "\n".join(f"{i+1}. {m}" for i, m in enumerate(user_msgs))
                    response += f"आपने पहले ये सवाल पूछे थे (Your previous questions):\n{numbered}"
                else:
                    response += "अभी तक कोई पिछला सवाल नहीं मिला। (No previous messages found in history.)"
            else:
                response += "मुझे अभी आपकी पिछली बातचीत की जानकारी नहीं है। (I don't have access to your chat history right now.)"
        else:
            # Context-aware fallback: use history if available, else generic help prompt
            if history:
                recent_user_msgs = [msg["content"] for msg in history if msg.get("role") == "user"]
                recent_bot_msgs  = [msg["content"] for msg in history if msg.get("role") == "assistant"]
                if recent_user_msgs:
                    last_q = recent_user_msgs[-1]
                    response += (
                        f"मैं आपकी बात समझ रहा हूँ। आपने पहले पूछा था: \"{last_q}\"\n\n"
                        "कृपया अधिक विवरण दें — जैसे फसल का नाम, समस्या या जानकारी जो आप चाहते हैं, "
                        "ताकि मैं आपको सटीक सलाह दे सकूँ।\n"
                        "(Please provide more details — crop name, problem or info needed — so I can give you precise advice.)"
                    )
                else:
                    response += (
                        "मैं आपकी किस विषय में मदद कर सकता हूँ?\n"
                        "कृपया बताएं: फसल का नाम, समस्या (जैसे रोग, सिंचाई, खाद) या बहीखाता दर्ज करना।"
                    )
            else:
                response += (
                    "नमस्ते! मैं आपका किसान साथी हूँ। मैं आपको फसल प्रबंधन, सिंचाई, और बहीखाता (ledger) में मदद कर सकता हूँ।\n"
                    "आप मुझसे गेहूं या आलू की खेती के बारे में सवाल पूछ सकते हैं, या कोई बिक्री दर्ज करने के लिए कह सकते हैं।"
                )
        
    for word in response.split(" "):
        yield word + " "
        time.sleep(0.04)

def generate_llama_cpp(prompt, system="", image_path=None, audio_path=None, history=None):
    """Query the in-process llama-cpp-python model with a timeout fallback to mock."""
    # 1. Fallback if llama_cpp is disabled (too slow or failed)
    if getattr(generate_llama_cpp, "disabled", False):
        print("[llm.py] llama_cpp is disabled (too slow or failed). Using mock backend.")
        for chunk in generate_mock(prompt, system, image_path, audio_path, history):
            yield chunk
        return

    # Acquire lock to prevent concurrent context corruption in llama.cpp
    acquired = _model_lock.acquire(blocking=True)
    if not acquired:
        print("[llm.py] Could not acquire model lock. Falling back to mock.")
        for chunk in generate_mock(prompt, system, image_path, audio_path, history):
            yield chunk
        return

    try:
        start_time = time.time()
        
        # Let's time how long it takes to get the model or if it hangs
        model = None
        try:
            model = init_llama_cpp()
        except Exception as e:
            print(f"[llm.py] Exception during init_llama_cpp: {e}")
            
        if model is None:
            print("[llm.py] Fallback to mock backend.")
            for chunk in generate_mock(prompt, system, image_path, audio_path, history):
                yield chunk
            return

        # Check if initialization took too long (exceeds 120s limit, indicating slow CPU)
        init_duration = time.time() - start_time
        if init_duration > 120.0:
            print(f"[llm.py] Warning: Model loading took {init_duration:.2f}s (exceeded 120s limit). Disabling llama_cpp and falling back to mock backend.")
            generate_llama_cpp.disabled = True
            for chunk in generate_mock(prompt, system, image_path, audio_path, history):
                yield chunk
            return

        # Prepend voice transcription if audio is supplied
        voice_prefix = ""
        if audio_path:
            transcription = transcribe_audio(audio_path, prompt)
            voice_prefix = f"[🎙️ **आवाज का अनुवाद (ASR):** \"{transcription}\"]\n\n"
            prompt = f"The user asked by voice: '{transcription}'. Response to this query."

        if image_path:
            # Note: Llama GGUF text-only files do not support direct image projection without mmproj.
            # We append a placeholder in the text so the LLM knows an image was submitted.
            prompt = f"[📸 छवि अपलोड की गई है / Image uploaded] {prompt}"

        formatted_prompt = f"<|im_start|>system\n{system}<|im_end|>\n"
        if history:
            for msg in history:
                role = msg.get("role", "user")
                content = msg.get("content", "")
                formatted_prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
        formatted_prompt += f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
        
        print(f"\n--- [llama.cpp INPUT PROMPT] ---\n{formatted_prompt}\n--------------------------------")
        print("--- [llama.cpp STREAMING RESPONSE] ---")
        try:
            response = model(
                formatted_prompt,
                max_tokens=512,
                temperature=0.3,
                top_p=0.9,
                stream=True
            )
            
            # Measure prompt prefill / first token time.
            # If it takes more than 120 seconds, fall back to mock.
            first_token_timeout = 120.0
            response_iter = iter(response)
            
            first_chunk_start = time.time()
            try:
                first_chunk = next(response_iter)
            except StopIteration:
                first_chunk = None
                
            prefill_duration = time.time() - first_chunk_start
            if prefill_duration > first_token_timeout:
                print(f"[llm.py] Prompt evaluation took {prefill_duration:.2f}s (exceeded {first_token_timeout}s limit). Disabling llama_cpp and falling back to mock.")
                generate_llama_cpp.disabled = True
                for chunk in generate_mock(prompt, system, image_path, audio_path, history):
                    yield chunk
                return

            if voice_prefix:
                yield voice_prefix
                
            if first_chunk:
                text = first_chunk['choices'][0]['text']
                cleaned = text.replace("<|im_end|>", "")
                print(cleaned, end="", flush=True)
                yield cleaned

            for chunk in response_iter:
                text = chunk['choices'][0]['text']
                cleaned = text.replace("<|im_end|>", "")
                print(cleaned, end="", flush=True)
                yield cleaned
            print("\n--------------------------------------")
        except Exception as e:
            print(f"[llm.py] Error running llama.cpp: {e}. Falling back to mock.")
            for chunk in generate_mock(prompt, system, image_path, audio_path, history):
                yield chunk
    finally:
        _model_lock.release()

def generate(prompt, system="", image_path=None, audio_path=None, history=None, stream=True):
    """Entry point for LLM generation supporting text, image, and voice inputs."""
    print(f"[llm.py] Using backend: {BACKEND}")
    if BACKEND == "llama_cpp":
        generator = generate_llama_cpp(prompt, system, image_path, audio_path, history)
    else: # mock
        generator = generate_mock(prompt, system, image_path, audio_path, history)
        
    if stream:
        return generator
    else:
        res = ""
        for chunk in generator:
            res += chunk
        return res