# Model loading utilities - Deployment Optimized

import streamlit as st
import whisper
import easyocr
import os
from sentence_transformers import SentenceTransformer
from transformers import CLIPProcessor, CLIPModel
from gpt4all import GPT4All
from config import MODELS_CONFIG, DEPLOYMENT_MODE
from device_setup import setup_device
@st.cache_resource
def load_models():
    try:
        # Writable caches first (prevent /.cache writes)
        import pathlib
        cache_root = os.environ.get("HF_HOME", "/tmp/hf")
        pathlib.Path(cache_root).mkdir(parents=True, exist_ok=True)

        # Make sure all relevant env vars point to the same place
        os.environ.setdefault("HF_HOME", cache_root)
        os.environ.setdefault("TRANSFORMERS_CACHE", cache_root)
        os.environ.setdefault("HUGGINGFACE_HUB_CACHE", cache_root)
        os.environ.setdefault("HF_DATASETS_CACHE", cache_root)
        os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", cache_root)
        os.environ.setdefault("XDG_CACHE_HOME", "/tmp")

        # Setup device
        device, has_gpu = setup_device()
        device_str = "cuda" if has_gpu else "cpu"

        progress = st.progress(0, text="Initializing models...")

        # Text embeddings (explicit cache folder)
        progress.progress(0.1, text="Loading text encoder...")
        st.write("📝 Loading text encoder...")
        text_encoder = SentenceTransformer(
            MODELS_CONFIG["text_encoder"],
            device=device,
            cache_folder=cache_root
        )

        # CLIP (route Hugging Face cache explicitly)
        progress.progress(0.3, text="Loading CLIP model...")
        st.write("🖼️ Loading CLIP model...")
        clip_model_name = MODELS_CONFIG["clip_model"]
        clip_model = CLIPModel.from_pretrained(clip_model_name, cache_dir=cache_root).to(device)
        clip_processor = CLIPProcessor.from_pretrained(clip_model_name, cache_dir=cache_root)

        # Whisper (force CPU on free tier unless GPU present)
        progress.progress(0.5, text="Loading Whisper...")
        st.write("🎤 Loading Whisper model...")
        whisper_size = MODELS_CONFIG["whisper_base"]
        whisper_model = whisper.load_model(whisper_size, device=device_str)

        # EasyOCR (avoid CUDA probing on CPU Spaces)
        progress.progress(0.7, text="Loading OCR...")
        st.write("🔍 Loading OCR reader...")
        ocr_reader = easyocr.Reader(['en'], gpu=False if not has_gpu else True)

        # Local LLM
        progress.progress(0.9, text="Loading LLM...")
        st.write("🧠 Loading local LLM...")
        try:
            llm = GPT4All(MODELS_CONFIG["llm_model"], allow_download=True)
            st.write("✅ LLM loaded successfully!")
        except Exception as llm_error:
            st.warning(f"⚠️ LLM loading failed: {llm_error}")
            st.info("💡 The system will work without local LLM for document search")
            llm = None

        progress.progress(1.0, text="All models loaded!")
        st.success("🎉 All models loaded successfully!")
        return text_encoder, clip_model, clip_processor, whisper_model, ocr_reader, llm, device, has_gpu

    except Exception as e:
        st.error(f"❌ Error loading models: {e}")
        st.info("💡 Try refreshing the page or check your internet connection")
        return None, None, None, None, None, None, None, False

def get_model_info():
    '''Get information about loaded models'''
    return {
        "text_encoder": MODELS_CONFIG["text_encoder"],
        "clip_model": MODELS_CONFIG["clip_model"], 
        "whisper_model": MODELS_CONFIG["whisper_base"],
        "llm_model": MODELS_CONFIG["llm_model"],
        "deployment_mode": DEPLOYMENT_MODE
    }