import os, gc, time, logging, threading
import numpy as np
import faiss
import gradio as gr
from sentence_transformers import SentenceTransformer
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

os.environ["HF_HOME"] = "/app/cache"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

MODEL_DIR      = "/app/cache/model"
INDEX_PATH     = "/app/cache/index.faiss"
TEXT_CACHE     = "/app/cache/chunks.npy"

MODEL_REPO     = "unsloth/Qwen3-1.7B-GGUF"
MODEL_FILENAME = "Qwen3-1.7B-Q4_K_M.gguf"
MODEL_PATH     = os.path.join(MODEL_DIR, MODEL_FILENAME)

EMBED_MODEL      = "BAAI/bge-small-en-v1.5"
BGE_QUERY_PREFIX = "Represent this sentence for searching relevant passages: "

N_CTX = 2048; N_THREADS = int(os.environ.get("N_THREADS", "2")); MAX_TOKENS = 150
os.makedirs(MODEL_DIR, exist_ok=True)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

llm = None; index = None; stored_texts = []; embedder = None; ready = False

def load_index():
    global index, stored_texts, embedder
    logger.info("🔧 Loading embedder...")
    embedder = SentenceTransformer(EMBED_MODEL, device="cpu")
    logger.info("🚀 Loading pre-built index...")
    index = faiss.read_index(INDEX_PATH)
    stored_texts = np.load(TEXT_CACHE, allow_pickle=True).tolist()
    logger.info(f"✅ {len(stored_texts)} chunks loaded!")

def retrieve(query, top_k=3):
    qvec = embedder.encode([BGE_QUERY_PREFIX + query], normalize_embeddings=True, convert_to_numpy=True)
    D, I = index.search(qvec.astype("float32"), top_k * 2)
    return [stored_texts[i] for s, i in zip(D[0], I[0]) if s > 0.4 and i < len(stored_texts)][:top_k]

def smart_retrieve(question):
    q_lower = question.lower()
    acts = {"wiba":"WORK_INJURY","work injury":"WORK_INJURY","compensation":"WORK_INJURY",
            "employer":"WORK_INJURY","noise":"Noise","sound":"Noise","fire":"Fire-Risk",
            "building":"Building Code","construction":"Building Code","waste":"WASTE",
            "wetland":"Wetlands","public health":"Public Health","toxic":"Toxic",
            "chemical":"Toxic","hazardous":"Toxic","environment":"Environment","evidence":"Evidence"}
    strategy = "broad_search"
    for kw, act in acts.items():
        if kw in q_lower: strategy = f"targeted:{act}"; break

    results = retrieve(question, top_k=3)
    if strategy.startswith("targeted:"):
        act = strategy.split(":")[1].lower()
        targeted = [r for r in results if act[:6] in r.lower()]
        if targeted: results = targeted[:2]

    ctx = "\n---\n".join(results) if results else ""
    if len(ctx) > 700: ctx = ctx[:700] + "..."
    return ctx, strategy

def init_llm():
    global llm
    if not os.path.exists(MODEL_PATH):
        logger.info(f"📥 Downloading {MODEL_FILENAME}...")
        hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME, local_dir=MODEL_DIR)
    logger.info("🚀 Loading Qwen3 1.7B...")
    llm = Llama(model_path=MODEL_PATH, n_ctx=N_CTX, n_threads=N_THREADS, n_batch=512, verbose=False)
    logger.info("✅ LLM ready.")

SYSTEM = "You are an OSH legal assistant. Answer in 2-3 sentences from Context. Cite Act/section."

def agentic_rag(question):
    if not ready: return "⏳ System loading — please retry."
    if not question.strip(): return "Please enter a question."

    t0 = time.time()
    ctx, strategy = smart_retrieve(question)
    t_retrieve = time.time() - t0
    if not ctx: return "No relevant documents found."

    t2 = time.time()
    # Pre-fill empty think block to skip Qwen3 thinking mode
    prompt = (
        f"<|im_start|>system\n{SYSTEM}<|im_end|>\n"
        f"<|im_start|>user\n/no_think\nContext:\n{ctx}\n\nQuestion: {question}<|im_end|>\n"
        f"<|im_start|>assistant\n<think>\n</think>\n"
    )
    output = llm(prompt, max_tokens=MAX_TOKENS, temperature=0.3, stop=["<|im_end|>", "<|im_start|>"])
    answer = output["choices"][0]["text"].strip()
    # Clean any leftover think tags
    answer = answer.replace("<think>", "").replace("</think>", "").strip()
    t_generate = time.time() - t2

    total = time.time() - t0
    return answer + f"\n\n---\n🤖 `{strategy}` · Retrieve {t_retrieve:.1f}s · Generate {t_generate:.1f}s · **Total {total:.1f}s**"

def startup():
    global ready
    load_index(); gc.collect()
    init_llm(); gc.collect()
    ready = True
    logger.info("🎉 Ready!")

threading.Thread(target=startup, daemon=True).start()

with gr.Blocks(title="OSH Agentic RAG", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 OSH Agentic RAG\n**Qwen3 1.7B** · BGE-small · Smart routing\n\nAsk about WIBA, Noise, Fire, Building Code, Public Health, Environment, Waste, Wetlands, Evidence Act.")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🧠 Smart Routing\nAuto-detects: WIBA, Noise, Fire, Building, Health, Environment, Waste, Wetlands, Evidence")
        with gr.Column(scale=2):
            q = gr.Textbox(label="Question", placeholder="What are noise exposure limits?", lines=2)
            btn = gr.Button("🤖 Ask", variant="primary", size="lg")
            out = gr.Markdown(label="Answer")
            btn.click(fn=agentic_rag, inputs=[q], outputs=[out], api_name="ask")
            q.submit(fn=agentic_rag, inputs=[q], outputs=[out], api_name=False)
    gr.Markdown("---\n```python\nfrom gradio_client import Client\nclient = Client('Rofati/osh-agentic-rag')\nresult = client.predict('What is WIBA?', api_name='/ask')\n```")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)