# app.py
import os
import glob
from huggingface_hub import snapshot_download
from ctransformers import AutoModelForCausalLM
import gradio as gr

# Model info
MODEL_REPO = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
MODEL_FILE = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
MODEL_CACHE_DIR = os.environ.get("MODEL_DIR", "/app/model_cache")

os.makedirs(MODEL_CACHE_DIR, exist_ok=True)

# If model not present, download it (public HF repo)
def ensure_model():
    # If user committed model locally, prefer that
    local_paths = glob.glob(os.path.join(MODEL_CACHE_DIR, "**", MODEL_FILE), recursive=True)
    candidate = local_paths[0] if local_paths else os.path.join(MODEL_CACHE_DIR, MODEL_FILE)
    if os.path.exists(candidate):
        return candidate

    print("Downloading model snapshot from Hugging Face hub...")
    repo_path = snapshot_download(repo_id=MODEL_REPO, cache_dir=MODEL_CACHE_DIR)
    matches = glob.glob(os.path.join(repo_path, "**", MODEL_FILE), recursive=True)
    if matches:
        return matches[0]
    # fallback: check repo_path root
    fallback = os.path.join(repo_path, MODEL_FILE)
    if os.path.exists(fallback):
        return fallback
    raise FileNotFoundError(f"Could not locate {MODEL_FILE} after download (repo_path={repo_path}).")

model_path = ensure_model()
print("Using model file:", model_path)

# Load model (ctransformers handles GGUF)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_REPO,
    model_file=model_path,
    model_type="llama",
    max_new_tokens=256,
    temperature=0.7
)

custom_data = {
    "hi": "hello Good morning,i am chatbot deepika",
    "about you": "I am deepika AI Chatbot,how may i help you..",
    "i love you": "i love you too",
    "age": "45"
}

def chatbot(msg: str):
    if not msg:
        return ""
    if msg in custom_data:
        return custom_data[msg]
    # ctransformers returns a string for simple calls
    return model(f"### Instruction:\n{msg}\n\n### Response:\n")

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    gr.Interface(fn=chatbot, inputs="text", outputs="text", title="Custom Chatbot - TinyLlama")\
      .launch(server_name="0.0.0.0", server_port=port)