# app.py — VERA backend (improved)
import os
import re
import uvicorn
import json
import time
import hashlib
import logging
import requests
from collections import deque

import uuid
from fastapi.staticfiles import StaticFiles
from fastapi import FastAPI, Request, Header
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware

from sentence_transformers import SentenceTransformer
from pinecone import Pinecone, ServerlessSpec

# ───────────────────────────────── Logging ─────────────────────────────────────
logging.basicConfig(level=logging.INFO)
log = logging.getLogger("vera-backend")

# ───────────────────────────────── Env / Config ───────────────────────────────
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
HF_API_KEY       = os.getenv("HF_API_KEY")

# realtime search (pick one you have a key for; Tavily recommended)
TAVILY_API_KEY   = os.getenv("TAVILY_API_KEY")  # https://www.tavily.com/
BING_API_KEY = os.getenv("BING_API_KEY")
# Optional: You can also add SERPAPI_API_KEY / BING_API_KEY if you prefer

URL_DECISION   = os.getenv("URL_DECISION",   "https://thevera-decision-making-model.hf.space/generate_text")
URL_SUMMARIZER = os.getenv("URL_SUMMARIZER", "https://thevera-get-best-answer.hf.space/summarize")

INSTR_DECISION = os.getenv("INSTR_DECISION", "instructions_L1.txt")
INSTR_HEALTH   = os.getenv("INSTR_HEALTH",   "instructions_health.txt")

REGION     = os.getenv("PINECONE_REGION", "us-east-1")
INDEX_NAME = os.getenv("PINECONE_INDEX",  "veradb")

EMBED_MODEL_ID = os.getenv("EMBED_MODEL_ID", "distiluse-base-multilingual-cased-v1")
EMBED_DIM      = int(os.getenv("EMBED_DIM", "512"))

# Memory / session config
SESSION_TTL_SEC   = int(os.getenv("SESSION_TTL_SEC", "7200"))  # 2 hours
SESSION_MAX_TURNS = int(os.getenv("SESSION_MAX_TURNS", "10"))  # keep last 10 items (user+assistant)

# Timeouts
HTTP_TIMEOUT = int(os.getenv("HTTP_TIMEOUT", "60"))

# ───────────────────────────────── FastAPI App ────────────────────────────────
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # or restrict to your static Space origin
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

os.makedirs("generated", exist_ok=True)
app.mount("/generated", StaticFiles(directory="generated"), name="generated")

# ───────────────────────────────── Dependencies ───────────────────────────────
# Pinecone (guard if key missing)
pc = None
index = None
if PINECONE_API_KEY:
    try:
        pc = Pinecone(api_key=PINECONE_API_KEY)
        if INDEX_NAME not in pc.list_indexes().names():
            pc.create_index(
                name=INDEX_NAME,
                dimension=EMBED_DIM,
                metric="cosine",
                spec=ServerlessSpec(cloud="aws", region=REGION),
            )
        index = pc.Index(INDEX_NAME)
        log.info("Pinecone index ready: %s", INDEX_NAME)
    except Exception:
        log.exception("Pinecone init failed")
else:
    log.warning("PINECONE_API_KEY not set — retrieval disabled.")

# Embeddings
embed_model = SentenceTransformer(EMBED_MODEL_ID)

# ───────────────────────────────── Utilities ──────────────────────────────────
def read_file(path: str) -> str:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        log.error("Failed to read %s: %s", path, e)
        return ""

def post_hf(url: str, payload: dict):
    """Call HF microservice with bearer if provided; return text or ''."""
    try:
        s = requests.Session()
        if HF_API_KEY:
            s.headers.update({"Authorization": f"Bearer {HF_API_KEY}"})
        r = s.post(url, json=payload, timeout=HTTP_TIMEOUT)
        if r.status_code != 200:
            log.error("HF call failed %s: %s", r.status_code, r.text[:240])
            return ""
        data = r.json()
        return data.get("response") or data.get("answer") or ""
    except Exception:
        log.exception("HF call exception")
        return ""

def detect_language(text: str) -> str:
    # Devanagari → 'hi' (Hindi/Sanskrit), else English
    return "hi" if any("\u0900" <= ch <= "\u097F" for ch in text) else "en"

def clamp(n, lo, hi):
    try:
        n = float(n)
    except Exception:
        return lo
    return max(lo, min(hi, n))

# ─────────────────────────────── In-memory sessions ───────────────────────────
# { session_id: {"history": deque[{"role": "user"/"assistant", "content": str}], "last_category": str, "t": float} }
SESSIONS = {}

def prune_sessions():
    now = time.time()
    stale = [sid for sid, s in SESSIONS.items() if (now - s.get("t", now)) > SESSION_TTL_SEC]
    for sid in stale:
        SESSIONS.pop(sid, None)

def get_session_id(request: Request, body: dict) -> str:
    sid = (body.get("session_id") or "").strip()
    if sid:
        return sid
    # derive from ip + ua
    ip = request.client.host if request.client else "0.0.0.0"
    ua = request.headers.get("user-agent", "")
    raw = f"{ip}|{ua}"[:256]
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:32]

def append_history(session_id: str, role: str, content: str):
    if not session_id:
        return
    prune_sessions()
    s = SESSIONS.setdefault(session_id, {"history": deque(maxlen=SESSION_MAX_TURNS*2), "last_category": None, "t": time.time()})
    s["history"].append({"role": role, "content": content})
    s["t"] = time.time()

def recent_history_text(session_id: str, max_chars: int = 4000) -> str:
    s = SESSIONS.get(session_id)
    if not s:
        return ""
    parts = []
    for item in list(s["history"])[-SESSION_MAX_TURNS*2:]:
        parts.append(f'{item["role"]}: {item["content"]}')
        if sum(len(p) for p in parts) > max_chars:
            break
    return "\n".join(parts)

def set_last_category(session_id: str, cat: str):
    if not session_id:
        return
    s = SESSIONS.setdefault(session_id, {"history": deque(maxlen=SESSION_MAX_TURNS*2), "last_category": None, "t": time.time()})
    s["last_category"] = cat
    s["t"] = time.time()

def get_last_category(session_id: str):
    s = SESSIONS.get(session_id)
    return s.get("last_category") if s else None

# ─────────────────────────────── Decision helpers ─────────────────────────────
def decide_category(user_text: str, session_id: str = "") -> str:
    """
    Returns exactly one of:
    health_wellness, spiritual_guidance, generate_image, realtime_query, other_query
    """
    # Quick heuristics (avoid a remote call for obvious cases)
    lt = user_text.strip().lower()
    if lt.startswith("/image") or lt.startswith("image:"):
        return "generate_image"
    if any(k in lt for k in ["latest", "today", "this week", "breaking", "price now", "live score"]):
        return "realtime_query"

    # Follow-up heuristics: short/anaphoric → reuse last category
    if len(lt) < 40 and re.search(r"\b(continue|and\?|more|what about|same|that|it|those)\b", lt):
        prev = get_last_category(session_id)
        if prev:
            return prev

    # Model decision
    instructions = read_file(INSTR_DECISION)
    payload = {"prompt": user_text, "instructions": instructions}
    raw = (post_hf(URL_DECISION, payload) or "").strip().lower()
    mapping = {
        "health_wellness": "health_wellness",
        "spiritual_guidance": "spiritual_guidance",
        "generate_image": "generate_image",
        "realtime_query": "realtime_query",
        "other_query": "other_query",
    }
    return mapping.get(raw, "other_query")

def decide_health_namespace(user_text: str) -> str:
    """
    Returns ayurvedic, western, or both (lowercase).
    """
    instructions = read_file(INSTR_HEALTH)
    payload = {"prompt": user_text, "instructions": instructions}
    raw = (post_hf(URL_DECISION, payload) or "").strip().lower()
    return raw if raw in {"ayurvedic", "western", "both"} else "both"

# ───────────────────────────── Pinecone retrieval ─────────────────────────────
def query_pinecone(namespace: str, prompt: str, top_k: int = 12):
    if not index:
        return []
    try:
        vec = embed_model.encode(prompt).tolist()
        res = index.query(namespace=namespace, vector=vec, top_k=top_k, include_metadata=True)
        return res.get("matches", []) or []
    except Exception:
        log.exception("Pinecone query failed for ns=%s", namespace)
        return []

def build_context(matches, cap_chars: int = 6000) -> str:
    ctx, size = [], 0
    for m in matches:
        meta = m.get("metadata") or {}
        t = (meta.get("text") or "").strip()
        if not t:
            continue
        ctx.append(t)
        size += len(t)
        if size >= cap_chars:
            break
    return "\n\n".join(ctx)

# ───────────────────────────── Realtime web search ────────────────────────────

def search_bing(query: str, k: int = 5):
    items = []
    if not BING_API_KEY:
        return items
    try:
        r = requests.get(
            "https://api.bing.microsoft.com/v7.0/search",
            headers={"Ocp-Apim-Subscription-Key": BING_API_KEY},
            params={"q": query, "mkt": "en-US", "count": k, "textDecorations": False, "answerCount": k},
            timeout=HTTP_TIMEOUT
        )
        if r.status_code == 200:
            j = r.json()
            for w in (j.get("webPages", {}).get("value", [])[:k]):
                items.append({
                    "title": w.get("name") or "",
                    "url": w.get("url") or "",
                    "content": w.get("snippet") or "",
                })
    except Exception:
        log.exception("Bing search failed")
    return items

    
def search_web(query: str, k: int = 5):
    """
    Returns list of {title, url, content}.
    Priority: Tavily -> DuckDuckGo IA -> Wikipedia summaries
    """
    results = []

    # 1) Tavily
    if TAVILY_API_KEY:
        try:
            r = requests.post(
                "https://api.tavily.com/search",
                json={"api_key": TAVILY_API_KEY, "query": query, "max_results": k, "include_answer": False},
                timeout=HTTP_TIMEOUT,
            )
            if r.status_code == 200:
                data = r.json()
                for item in (data.get("results") or [])[:k]:
                    results.append({
                        "title": item.get("title") or "",
                        "url": item.get("url") or "",
                        "content": item.get("content") or item.get("snippet") or "",
                    })
        except Exception:
            log.exception("Tavily search failed")

    # 2) Bing fallback
    if not results:
        results = search_bing(query, k=k)

    # 3) DuckDuckGo Instant Answer (no key)
    if not results:
        try:
            r = requests.get(
                "https://api.duckduckgo.com/",
                params={"q": query, "format": "json", "no_html": 1, "no_redirect": 1},
                timeout=HTTP_TIMEOUT,
            )
            if r.status_code == 200:
                data = r.json()
                if data.get("AbstractText"):
                    results.append({
                        "title": data.get("Heading") or "",
                        "url": data.get("AbstractURL") or "",
                        "content": data.get("AbstractText") or "",
                    })
                for rt in (data.get("RelatedTopics") or []):
                    if isinstance(rt, dict) and rt.get("Text"):
                        results.append({
                            "title": rt.get("Text")[:80],
                            "url": (rt.get("FirstURL") or ""),
                            "content": rt.get("Text") or "",
                        })
                results = results[:k]
        except Exception:
            log.exception("DuckDuckGo IA failed")

    # 4) Wikipedia (simple)
    if not results:
        try:
            r = requests.get("https://en.wikipedia.org/w/api.php",
                             params={"action":"opensearch","search":query,"limit":k,"namespace":0,"format":"json"},
                             timeout=HTTP_TIMEOUT)
            if r.status_code == 200:
                data = r.json()
                titles = data[1] or []
                urls   = data[3] or []
                for t,u in list(zip(titles, urls))[:k]:
                    results.append({"title": t, "url": u, "content": t})
        except Exception:
            log.exception("Wikipedia search failed")

    return results[:k]

def summarize_search(query: str, lang: str) -> str:
    items = search_web(query, k=5)
    if not items:
        return "I couldn't find current information right now. Please try rephrasing or ask a more specific query."

    # Compose context for the summarizer
    ctx_parts = []
    for it in items:
        ctx_parts.append(f"{it['title']}\n{it['content']}\nSource: {it['url']}\n")
    context = "\n\n".join(ctx_parts)

    payload = {
        "context": context,
        "question": f"Using the sources above, answer this query concisely and include a friendly tone.\nQuery: {query}",
        "language": "Hindi" if lang == "hi" else "English",
        "reply_type": "brief",
    }
    ans = (post_hf(URL_SUMMARIZER, payload) or "").strip()

    # Add sources footer
    src_lines = []
    for i, it in enumerate(items[:3], 1):
        if not it.get("url"):
            continue
        src_lines.append(f"{i}. {it['title'] or it['url']} — {it['url']}")
    if src_lines:
        ans += "\n\nSources:\n" + "\n".join(src_lines)

    return ans or "No answer available."

# ───────────────────────────── Domain handlers ────────────────────────────────
def handle_spiritual(user_text: str, lang: str) -> str:
    # Try existing namespace "spritual" then sane "spiritual"
    matches = query_pinecone("spritual", user_text, top_k=20)
    if not matches:
        matches = query_pinecone("spiritual", user_text, top_k=20)

    if not matches:
        # Positive fallback
        base = "Hindi" if lang == "hi" else "English"
        payload = {
            "context": "",
            "question": f"Provide compassionate spiritual guidance. Question: {user_text}",
            "language": base,
            "reply_type": "detailed",
        }
        return (post_hf(URL_SUMMARIZER, payload) or "Let's reflect on this with kindness and clarity.").strip()

    context = build_context(matches)
    reply_type = "brief" if re.search(r"\b(brief|short|tl;dr)\b", user_text.lower()) else "detailed"
    payload = {
        "context": context,
        "question": user_text,
        "language": "Hindi" if (lang == "hi" and "in english" not in user_text.lower()) else "English",
        "reply_type": reply_type,
    }
    ans = post_hf(URL_SUMMARIZER, payload) or ""
    return ans.strip() or "May this perspective be helpful."

def merge_health(ay, we):
    merged = (ay or []) + (we or [])
    merged.sort(key=lambda m: m.get("score", 0), reverse=True)
    return merged[:12]

def handle_health(user_text: str, lang: str, image_text: str = "") -> str:
    choice = decide_health_namespace(user_text)
    if choice == "ayurvedic":
        matches = query_pinecone("ayurvedic", user_text, 14)
    elif choice == "western":
        matches = query_pinecone("western", user_text, 14)
    else:
        ay = query_pinecone("ayurvedic", user_text, 10)
        we = query_pinecone("western", user_text, 10)
        matches = merge_health(ay, we)

    base_lang = "Hindi" if lang == "hi" else "English"
    reply_type = "detailed" if any(k in user_text.lower() for k in ["detail", "explain", "why", "how"]) else "brief"

    context = build_context(matches) if matches else ""
    if image_text:
        context = (context + "\n\nExtracted from health image:\n" + image_text).strip()

    payload = {
        "context": context,
        "question": user_text,
        "language": base_lang,
        "reply_type": reply_type,
    }
    ans = (post_hf(URL_SUMMARIZER, payload) or "").strip()
    disclaimer = ("\n\n**Note:** This is educational, not a medical diagnosis. "
                  "Consult a clinician for personalized care.")
    return ans + disclaimer if ans else "I couldn't parse enough context—could you add a few details?"


def handle_realtime(user_text: str, lang: str) -> str:
    # Real search + summarization
    return summarize_search(user_text, lang)

def handle_other(user_text: str, lang: str, context_extra: str = "") -> str:
    payload = {
        "context": context_extra or "",
        "question": user_text,
        "language": "Hindi" if lang == "hi" else "English",
        "reply_type": "detailed",
    }
    ans = post_hf(URL_SUMMARIZER, payload) or ""
    return ans.strip() or "Here’s a helpful overview."

def generate_image_free(prompt: str) -> str:
    """Return a URL to an image for the prompt. Prefers HF Inference -> fallback Pollinations."""
    # 1) Hugging Face Inference API (binary image)
    if HF_API_KEY:
        try:
            r = requests.post(
                "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1",
                headers={
                    "Authorization": f"Bearer {HF_API_KEY}",
                    "Accept": "image/png"
                },
                json={"inputs": prompt},
                timeout=120
            )
            if r.status_code == 200 and r.content:
                fname = f"{uuid.uuid4().hex}.png"
                path = os.path.join("generated", fname)
                with open(path, "wb") as f:
                    f.write(r.content)
                return f"/generated/{fname}"
            else:
                log.warning("HF image gen failed %s: %s", r.status_code, r.text[:200])
        except Exception:
            log.exception("HF image gen exception")

    # 2) Pollinations (no key; public)
    try:
        from urllib.parse import quote
        return f"https://image.pollinations.ai/prompt/{quote(prompt)}"
    except Exception:
        return ""


def handle_image(user_text: str, lang: str) -> str:
    prompt = user_text.replace("/image", "").strip() or user_text
    url = generate_image_free(prompt)
    if url:
        return f"Here’s your image for: **{prompt}**\n\n![{prompt}]({url})"
    return f"I couldn’t generate the image right now. Try again in a moment.\n\n**Prompt:** {prompt}"


# ───────────────────────────── Routes ─────────────────────────────────────────
@app.get("/")
def root():
    return {"ok": True, "service": "vera-backend"}

@app.get("/healthz")
def healthz():
    return {"ok": True}

@app.post("/reset_session")
async def reset_session(request: Request):
    body = await request.json()
    sid = (body.get("session_id") or "").strip()
    if sid and sid in SESSIONS:
        SESSIONS.pop(sid, None)
        return {"ok": True, "reset": True}
    return {"ok": True, "reset": False}

@app.post("/generate_result")
async def generate_result(request: Request, user_agent: str = Header(default="")):
    body = await request.json()
    raw_text = (body.get("message") or "").strip()
    image_text = (body.get("image_text") or "").strip()
    if not raw_text:
        return JSONResponse(status_code=400, content={"response": "No message provided"})

    session_id = get_session_id(request, body)

    # Parse a forced category from the header line
    allowed = {"health_wellness","spiritual_guidance","generate_image","realtime_query","other_query"}
    forcible = None
    text = raw_text
    if text.startswith("::category="):
        m = re.match(r"^::category\s*=\s*([a-zA-Z_]+)(?:\s+|[\r\n]+)?(.*)$", text, flags=re.S)
        if m:
            key = m.group(1).strip().lower()
            remainder = (m.group(2) or "").strip()
            if key in allowed:
                forcible = key
                text = remainder

    # Also allow JSON override
    force_from_body = (body.get("force_category") or "").strip().lower() or None
    if not forcible and force_from_body in allowed:
        forcible = force_from_body

    effective_text = text or raw_text
    lang = detect_language(effective_text)

    # Decide category (with session-aware heuristics)
    decision = forcible or decide_category(effective_text, session_id=session_id)
    set_last_category(session_id, decision)

    # Keep memory
    append_history(session_id, "user", effective_text)

    # Build optional conversation context for generic answers
    convo_ctx = recent_history_text(session_id)

    handlers = {
        "spiritual_guidance": lambda t, l: handle_spiritual(t, l),
        "health_wellness":   lambda t, l: handle_health(t, l, image_text=image_text),
        "realtime_query":    lambda t, l: handle_realtime(t, l),
        "generate_image":    lambda t, l: handle_image(t, l),
        "other_query":       lambda t, l: handle_other(t, l, context_extra=convo_ctx),
    }
    func = handlers.get(decision, handlers["other_query"])

    try:
        answer = func(effective_text, lang)
        append_history(session_id, "assistant", answer)
        return JSONResponse(content={"response": f"{decision}\n{answer}", "language": lang, "session_id": session_id})
    except Exception as e:
        log.exception("Handler failed")
        return JSONResponse(
            status_code=500,
            content={"response": f"other_query\nAn error occurred, but here's a quick tip: please try again in a moment.", "language": lang, "session_id": session_id}
        )

# ───────────────────────────── Main ───────────────────────────────────────────
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))