import os
import gradio as gr
from sentence_transformers import SentenceTransformer
from pathlib import Path
import torch
import numpy as np
from fastapi import FastAPI
import uvicorn
import shutil

# =====================================================
# 🧹 CACHE & STORAGE MANAGEMENT
# =====================================================
# Use temporary directories for caching to prevent disk overuse
os.environ["TRANSFORMERS_CACHE"] = "/tmp"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp"

# Optional: clear any pre-existing cache directories
shutil.rmtree(os.path.expanduser("~/.cache/torch"), ignore_errors=True)
shutil.rmtree(os.path.expanduser("~/.cache/huggingface"), ignore_errors=True)

# =====================================================
# 📁 MODEL SETUP
# =====================================================
# Instead of a heavy local snapshot, load the lightweight MiniLM model directly
MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# =====================================================
# 🚀 LOAD MODEL
# =====================================================
try:
    print(f"📁 Loading model: {MODEL_NAME}")
    model = SentenceTransformer(MODEL_NAME, device=DEVICE)
    print(f"✅ Model loaded successfully on {DEVICE.upper()}")
except Exception as e:
    raise RuntimeError(
        f"❌ Model loading failed. Ensure internet access or cached files are available.\nError: {e}"
    )

# =====================================================
# 🧠 EMBEDDING FUNCTION (for local + API use)
# =====================================================
def encode_text(text: str):
    """
    Returns a summary of the embedding vector (for Gradio interface).
    """
    if not text or not text.strip():
        return {"error": "Please enter some text."}

    try:
        embedding = model.encode([text])[0]
        preview = np.round(embedding[:10], 4).tolist()
        return {
            "embedding_shape": len(embedding),
            "embedding_preview (first 10 values)": preview,
        }
    except Exception as e:
        return {"error": f"Encoding failed: {str(e)}"}

# =====================================================
# ⚙️ API FUNCTION (for remote prediction)
# =====================================================
def api_predict(text: str):
    """
    Returns the full embedding vector as a list (for API use).
    """
    if not text or not text.strip():
        return []
    try:
        embedding = model.encode([text])[0]
        return embedding.tolist()
    except Exception as e:
        return {"error": str(e)}

# =====================================================
# 🎨 GRADIO INTERFACE
# =====================================================
demo = gr.Interface(
    fn=encode_text,
    inputs=gr.Textbox(label="Enter Text", placeholder="Type a sentence to encode..."),
    outputs=gr.JSON(label="Embedding Info"),
    title="💬 Sentence Embedding Generator",
    description=(
        "Encodes sentences into numerical embeddings using "
        "<b>MiniLM-L6-v2</b>. "
        "Useful for semantic similarity, clustering, or grouping."
    ),
    examples=[
        ["I love machine learning and natural language processing."],
        ["The movie was terrible, I wouldn't recommend it."],
        ["Artificial intelligence is transforming the world."],
    ],
)

# =====================================================
# 🚀 FASTAPI APP + GRADIO MOUNT
# =====================================================
app = FastAPI(title="Sentence Embedding API")

# Add the /predict endpoint
@app.post("/predict")
async def predict_api(data: dict):
    text = data.get("text", "")
    return api_predict(text)

# Mount Gradio UI at root "/"
app = gr.mount_gradio_app(app, demo, path="/")

# =====================================================
# 🔥 LAUNCH SERVER
# =====================================================
if __name__ == "__main__":
    print("🚀 Starting API and Gradio UI server...")
    uvicorn.run(app, host="0.0.0.0", port=7860)












# import gradio as gr
# from sentence_transformers import SentenceTransformer
# from pathlib import Path
# import torch
# import numpy as np

# # =====================================================
# # 📁 MODEL SETUP
# # =====================================================
# MODEL_DIR = Path(__file__).parent / "models" / "all-MiniLM" / "snapshots" / "c9745ed1d9f207416be6d2e6f8de32d1f16199bf"


# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# # =====================================================
# # 🚀 LOAD MODEL
# # =====================================================
# try:
#     print(f"📁 Loading model from: {MODEL_DIR}")
#     model = SentenceTransformer(str(MODEL_DIR), device=DEVICE)
#     print(f"✅ Model loaded successfully on {DEVICE.upper()}")
# except Exception as e:
#     raise RuntimeError(
#         f"❌ Model loading failed. Ensure the path '{MODEL_DIR}' "
#         f"is correct and contains model files. Original error: {e}"
#     )

# # =====================================================
# # 🧠 EMBEDDING FUNCTION (for local + API use)
# # =====================================================
# def encode_text(text: str):
#     """
#     Returns the embedding vector of a given text.
#     For API use, only returns the vector (list of floats).
#     """
#     if not text or not text.strip():
#         return {"error": "Please enter some text."}

#     try:
#         embedding = model.encode([text])[0]
#         preview = np.round(embedding[:10], 4).tolist()
#         return {
#             "embedding_shape": len(embedding),
#             "embedding_preview (first 10 values)": preview,
#         }
#     except Exception as e:
#         return {"error": f"Encoding failed: {str(e)}"}

# # =====================================================
# # ⚙️ API FUNCTION (for remote prediction)
# # =====================================================
# def api_predict(text: str):
#     """
#     This is the function that will be exposed via /predict API.
#     It returns only the embedding vector (as a list).
#     """
#     if not text or not text.strip():
#         return []
#     try:
#         embedding = model.encode([text])[0]
#         return embedding.tolist()
#     except Exception as e:
#         return {"error": str(e)}

# # =====================================================
# # 🎨 GRADIO INTERFACE
# # =====================================================
# demo = gr.Interface(
#     fn=encode_text,
#     inputs=gr.Textbox(label="Enter Text", placeholder="Type a sentence to encode..."),
#     outputs=gr.JSON(label="Embedding Info"),
#     title="💬 Sentence Embedding Generator",
#     description=(
#         "Encodes sentences into numerical embeddings using "
#         "<b>distilbert-base-nli-mean-tokens</b>. "
#         "Useful for semantic similarity, clustering, or grouping."
#     ),
#     examples=[
#         ["I love machine learning and natural language processing."],
#         ["The movie was terrible, I wouldn't recommend it."],
#         ["Artificial intelligence is transforming the world."],
#     ],
# )

# # =====================================================
# # 🧩 ADD API ENDPOINT (/predict)
# # =====================================================
# demo.add_api_route("/predict", api_predict, methods=["POST"])

# # =====================================================
# # 🚀 LAUNCH SERVER
# # =====================================================
# if __name__ == "__main__":
#     print("🚀 Starting Gradio API server...")
#     demo.launch(server_name="0.0.0.0", server_port=7860)