# app.py
import os
import uuid
import gradio as gr
from PIL import Image
from datetime import datetime
from supabase import create_client
from sentence_transformers import SentenceTransformer

# =======================
# SUPABASE CONFIG
# =======================
SUPABASE_URL = os.environ["https://pjrzisurkyeunhgbmvco.supabase.co"] # suprabase url
SUPABASE_KEY = os.environ["sb_publishable_mYlT4zRks_3BF-Lqqw35og_wzvLnMg2"] # suprabase service key
supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

# Table names
CONTEXT_TABLE = "context_memory"
CHAT_TABLE = "conversation_history"

# =======================
# USER SESSION MANAGEMENT
# =======================
def get_or_create_user_id(user_id_state):
    if user_id_state is None:
        return str(uuid.uuid4())
    return user_id_state

# =======================
# EMBEDDINGS SETUP
# =======================
embedding_model = SentenceTransformer(
    "jinaai/jina-embeddings-v3",
    trust_remote_code=True
)

TRUNCATE_DIM = 256  # matches your Supabase vector column

def embed_query(text: str):
    return embedding_model.encode([text], task="retrieval.query", truncate_dim=TRUNCATE_DIM)[0]

# =======================
# RAG RETRIEVAL
# =======================
TOP_K = 3

def get_top_k_chunks(user_id: str, query: str, k: int = TOP_K):
    q_emb = embed_query(query)

    # call Supabase RPC or SQL for vector similarity search
    # assuming you have a match_context_memory RPC defined in Supabase
    result = supabase.rpc(
        "match_context_memory",
        {
            "query_embedding": q_emb,
            "match_count": k,
            "user_id": user_id
        }
    ).execute()

    if result.data:
        return [row["semantic_chunk"] for row in result.data]
    return []

# =======================
# LOAD RECENT CHAT
# =======================
MAX_HISTORY = 10

def get_recent_history(user_id: str, limit: int = MAX_HISTORY):
    res = supabase.table(CHAT_TABLE) \
        .select("query_input,llm_output") \
        .eq("unique_user_id", user_id) \
        .order("convo_num", desc=True) \
        .limit(limit) \
        .execute()
    
    history = []
    for row in reversed(res.data):
        history.append(f"User: {row['query_input']}")
        if row["llm_output"]:
            history.append(f"Assistant: {row['llm_output']}")
    return "\n".join(history)

# =======================
# SAVE / UPDATE CHAT
# =======================
def insert_user_query(user_id, user_name, region, query):
    supabase.table(CHAT_TABLE).insert({
        "unique_user_id": user_id,
        "user_name": user_name,
        "region": region,
        "query_input": query
    }).execute()

def update_llm_output(user_id, llm_output):
    supabase.table(CHAT_TABLE) \
        .update({"llm_output": llm_output}) \
        .eq("unique_user_id", user_id) \
        .is_("llm_output", None) \
        .order("convo_num", desc=True) \
        .limit(1) \
        .execute()


# =======================
# VISION + LLM HANDLER (Bangla LLaMA 3.2)
# =======================
import requests

HF_LLM_KEY = os.environ.get("HF_API_KEY")
MODEL_ID = "hishab/titulm-llama-3.2-3b-v2.0"
HEADERS = {"Authorization": f"Bearer {HF_LLM_KEY}"}

def vision_llm(image: Image.Image | None, text: str, top_chunks: list, chat_history: str):
    """
    Calls Bangla LLaMA 3.2 via Hugging Face API with:
    - top-k RAG chunks
    - recent chat history
    - user input
    - optional image info
    """
    prompt_parts = [
        "You are a helpful assistant that can analyze images and text.",
        "Follow these rules: answer concisely, prefer facts from context.",
        "\n--- Context chunks ---\n" + "\n".join(top_chunks) if top_chunks else "",
        "\n--- Chat History ---\n" + chat_history if chat_history else "",
        "\n--- User Input ---\n" + text
    ]

    if image:
        prompt_parts.append(f"\nUser uploaded an image of size {image.width}x{image.height}.")

    prompt = "\n".join(prompt_parts)

    # Call Hugging Face Inference API
    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 512}}
    try:
        response = requests.post(
            f"https://api-inference.huggingface.co/models/{MODEL_ID}",
            headers=HEADERS,
            json=payload,
            timeout=60
        )
        response.raise_for_status()
        data = response.json()
        # HF API returns a list of dicts with "generated_text"
        if isinstance(data, list) and "generated_text" in data[0]:
            return data[0]["generated_text"]
        elif isinstance(data, dict) and "generated_text" in data:
            return data["generated_text"]
        else:
            return str(data)
    except Exception as e:
        return f"Error calling Bangla LLaMA 3.2: {e}"


# =======================
# GRADIO HANDLER
# =======================
def analyze_input(message, history, user_id_state):
    user_id = get_or_create_user_id(user_id_state)
    text = (message.get("text") or "").strip()
    files = message.get("files") or []
    
    image = None
    if files:
        image_path = files[0]["path"]
        image = Image.open(image_path)
    else:
        image_path = None

    # Save user query in Supabase
    insert_user_query(user_id, "Anonymous", "Unknown", text)

    # Retrieve top-k chunks for RAG
    top_chunks = get_top_k_chunks(user_id, text, k=TOP_K)
    
    # Load recent chat
    recent_history = get_recent_history(user_id, limit=MAX_HISTORY)

    # Call LLM
    llm_reply = vision_llm(image, text, top_chunks, recent_history)

    # Update chat with LLM reply
    update_llm_output(user_id, llm_reply)

    return llm_reply, user_id

# =======================
# GRADIO UI
# =======================
with gr.Blocks() as demo:
    user_id_state = gr.State(None)
    chatbot = gr.ChatInterface(
        fn=analyze_input,
        multimodal=True,
        additional_inputs=[user_id_state],
        additional_outputs=[user_id_state],
        title="InfoJury – Vision + RAG + Chat",
        description="Upload an image, ask a question, and the system will reason using your personal knowledge base plus vision."
    )

if __name__ == "__main__":
    demo.launch()