app/main.py

import logging

from fastapi import FastAPI, HTTPException, Header, Request
from fastapi.responses import RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import os

from app.api.schemas import ChatRequest, ChatResponse, HealthResponse, SessionUpdateNameRequest, SessionInitRequest, RateRequest, PollRequest, ResetRequest, DeleteChatRequest
from app.core.config import get_settings
from app.services.embeddings import EmbeddingService
from app.services.llm import LLMService
from app.services.rag_pipeline import RAGPipeline
from app.services.rate_limiter import InMemoryRateLimiter
from app.services.vector_store import FaissVectorStore
from app.services.reranker import RerankerService
from app.ui_gradio import demo
from app.admin.router import admin_router
from app.services import session_store as _ss
import asyncio
import gradio as gr

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

settings = get_settings()
embedding_service = EmbeddingService(settings.embedding_model)
vector_store = FaissVectorStore(
    embedding_service=embedding_service,
    docs_dir=settings.docs_dir,
    index_dir=settings.index_dir,
    chunk_size_tokens=settings.chunk_size_tokens,
    chunk_overlap_tokens=settings.chunk_overlap_tokens,
)
llm_service = LLMService(
    provider=settings.llm_provider,
    openai_api_key=settings.openai_api_key,
    openai_model=settings.openai_model,
    openai_rewrite_model=settings.openai_rewrite_model,
    groq_api_key=settings.groq_api_key,
    groq_model=settings.groq_model,
    groq_rewrite_model=settings.groq_rewrite_model,
    hf_api_key=settings.hf_api_key,
    hf_model=settings.hf_model,
    fireworks_api_key=settings.fireworks_api_key,
    fireworks_model=settings.fireworks_model,
    fireworks_rewrite_model=settings.fireworks_rewrite_model,
    timeout_s=settings.request_timeout_s,
)
reranker_service = RerankerService(settings.reranker_model)
pipeline = RAGPipeline(
    vector_store=vector_store, 
    llm_service=llm_service, 
    reranker=reranker_service,
    top_k=settings.top_k,
    max_context_chunks=settings.max_context_chunks
)

# ── FastAPI App ───────────────────────────────────────────────────────────────
fastapi_app = FastAPI(title=settings.app_name)

allow_origins = [o.strip() for o in settings.cors_allow_origins.split(",") if o.strip()]
fastapi_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], # Explicitly allow all for testing
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

rate_limiter = InMemoryRateLimiter(settings.rate_limit_requests, settings.rate_limit_window_seconds)

# ── Admin Panel (mounted on FastAPI before Gradio wrapping) ───────────────────
fastapi_app.include_router(admin_router)

# ── Static Files ──────────────────────────────────────────────────────────────
static_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "static")
if os.path.exists(static_dir):
    fastapi_app.mount("/static", StaticFiles(directory=static_dir), name="static")

@fastapi_app.get("/widget", include_in_schema=False)
async def get_widget():
    """Serves the chat widget HTML with explicit CORS headers."""
    from fastapi.responses import FileResponse
    path = os.path.join(static_dir, "addon.html")
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="Widget not found")
    return FileResponse(
        path, 
        headers={
            "Access-Control-Allow-Origin": "*",
            "Cache-Control": "no-cache, no-store, must-revalidate"
        }
    )

@fastapi_app.get("/sw.js", include_in_schema=False)
async def get_service_worker():
    """Serves the Service Worker at root scope."""
    from fastapi.responses import FileResponse
    path = os.path.join(static_dir, "sw.js")
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="Service worker not found")
    return FileResponse(
        path,
        media_type="application/javascript",
        headers={
            "Cache-Control": "no-cache, no-store, must-revalidate",
            "Service-Worker-Allowed": "/"
        }
    )

@fastapi_app.get("/martech_sol_logo.jpg", include_in_schema=False)
async def get_logo():
    """Serves the brand logo."""
    from fastapi.responses import FileResponse
    path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "martech_sol_logo.jpg")
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="Logo not found")
    return FileResponse(path, media_type="image/jpeg")


@fastapi_app.on_event("startup")
def startup_event() -> None:
    logger.info("─── RAG SERVICE STARTUP ───")
    logger.info("Provider: %s", settings.llm_provider)
    logger.info("Fireworks Key: %s", "PRESENT" if settings.fireworks_api_key else "MISSING")
    logger.info("Groq Key: %s", "PRESENT" if settings.groq_api_key else "MISSING")
    logger.info("Index Directory: %s", settings.index_dir)
    
    # Verify persistence mount
    if str(settings.index_dir).startswith("/data"):
        if os.path.exists("/data"):
            logger.info("✅ Persistent storage /data detected and in use.")
        else:
            logger.warning("⚠️ /data requested but NOT found on filesystem!")
    else:
        logger.warning("⚠️ Using ephemeral storage (data will be lost on rebuild).")

    logger.info("Loading/building FAISS index...")
    vector_store.build_or_load()
    logger.info("RAG service started with %s chunks", len(vector_store.metadata))


@fastapi_app.on_event("shutdown")
async def shutdown_event() -> None:
    """Gracefully close the shared httpx client to free connections."""
    logger.info("─── RAG SERVICE SHUTDOWN — closing HTTP client ───")
    await llm_service.close()


@fastapi_app.get("/", include_in_schema=False)
async def root():
    """Prevent direct browser access to the root URL."""
    raise HTTPException(
        status_code=403, 
        detail="Direct access to this resource is restricted. Please use the authorized interface."
    )


@fastapi_app.get("/health", response_model=HealthResponse)
def health() -> HealthResponse:
    h = vector_store.health()
    return HealthResponse(status="ok", **h)

async def send_whatsapp_message(text: str):
    """Sends a WhatsApp notification using settings from admin panel (persisted JSON) or env vars."""
    try:
        from app.admin.router import get_whatsapp_settings
        ws = get_whatsapp_settings()
        
        phone_id = ws.get("phone_id", "")
        token = ws.get("token", "")
        admin_number = ws.get("admin_number", "")
        
        if not phone_id or not token or not admin_number:
            logger.warning(
                "WhatsApp settings incomplete — phone_id: %s, token: %s, admin_number: %s. Skipping.",
                "SET" if phone_id else "MISSING",
                "SET" if token else "MISSING",
                "SET" if admin_number else "MISSING",
            )
            return
        
        admin_numbers = [n.strip() for n in admin_number.split(",") if n.strip()]
        if not admin_numbers:
            logger.warning("No valid WhatsApp numbers found after parsing, skipping.")
            return

        url = f"https://graph.facebook.com/v17.0/{phone_id}/messages"
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }
        
        logger.info("Sending WhatsApp notification to %d number(s)...", len(admin_numbers))
        
        def _send_sync():
            import urllib.request
            import urllib.error
            import json
            import ssl
            
            ctx = ssl.create_default_context()
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
            
            for number in admin_numbers:
                payload = {
                    "messaging_product": "whatsapp",
                    "to": number,
                    "type": "text",
                    "text": {"body": text}
                }
                data = json.dumps(payload).encode('utf-8')
                req = urllib.request.Request(url, data=data, headers=headers, method='POST')
                try:
                    with urllib.request.urlopen(req, timeout=15.0, context=ctx) as response:
                        logger.info("WhatsApp notification sent successfully to %s", number)
                except urllib.error.HTTPError as e:
                    err_msg = e.read().decode('utf-8')
                    logger.error("WhatsApp API error for %s (%d): %s", number, e.code, err_msg)
                except Exception as e:
                    logger.error("Network error sending WhatsApp to %s: %s - %s", number, type(e).__name__, str(e))
        
        import asyncio
        await asyncio.to_thread(_send_sync)
    except Exception as e:
        logger.error("Failed to send WhatsApp messages: %s", e, exc_info=True)


@fastapi_app.post("/api/chat", response_model=ChatResponse)
async def chat(
    payload: ChatRequest,
    request: Request,
    x_api_key: str | None = Header(default=None),
) -> ChatResponse:
    if not payload.message.strip():
        raise HTTPException(status_code=400, detail="message must not be empty")
    if settings.api_key and x_api_key != settings.api_key:
        raise HTTPException(status_code=401, detail="unauthorized")

    client_ip = request.client.host if request.client else "unknown"
    if not rate_limiter.allow(client_ip):
        raise HTTPException(status_code=429, detail="rate limit exceeded")

    logger.info("Incoming chat request: msg='%s', sid='%s'", payload.message[:30], payload.session_id)

    # Fetch user_name and handle session status
    user_name = None
    if payload.session_id:
        sess_data = await _ss.get_session(payload.session_id)
        if sess_data:
            user_name = sess_data.get("user_name")
            if sess_data.get("blocked"):
                raise HTTPException(status_code=403, detail="Forbidden. User is blocked.")
            status = sess_data.get("status", "AI")
            msg_lower = payload.message.lower()

            if status == "HUMAN_TAKEOVER":
                asyncio.create_task(_ss.save_message(payload.session_id, payload.message, ""))
                return ChatResponse(reply="", retrieved_chunks=[])
                
            if status == "QUEUED":
                reply_msg = "Please wait, an agent will be with you shortly."
                asyncio.create_task(_ss.save_message(payload.session_id, payload.message, reply_msg))
                return ChatResponse(reply=reply_msg, retrieved_chunks=[])
                
            if status == "CONFIRMING_TRANSFER":
                if msg_lower in ["yes", "y", "yeah", "yep", "sure", "ok", "okay"]:
                    await _ss.update_session_field(payload.session_id, "status", "QUEUED")
                    reply_msg = "You are in the queue and the HR/supervising authority will contact you soon."
                    asyncio.create_task(_ss.save_message(payload.session_id, payload.message, reply_msg))
                    asyncio.create_task(send_whatsapp_message(
                        f"🚨 Human Handoff Request\n\n"
                        f"👤 User: {user_name or 'Unknown'}\n"
                        # f"📋 Session: {payload.session_id}\n\n"
                        f"{user_name or 'A user'} wants to talk to the HR/supervising authority.\n"
                        f"Please check the Admin Panel to respond."
                    ))
                    return ChatResponse(reply=reply_msg, retrieved_chunks=[])
                else:
                    await _ss.update_session_field(payload.session_id, "status", "AI")
                    reply_msg = "Transfer cancelled. How can I help you?"
                    asyncio.create_task(_ss.save_message(payload.session_id, payload.message, reply_msg))
                    return ChatResponse(reply=reply_msg, retrieved_chunks=[])


    try:
        result = await pipeline.chat(payload.message, payload.history, user_name=user_name)
        reply_text = result["reply"]
        
        # ── AI Intent Handover Detection ──
        if "<HUMAN_TAKEOVER_REQUEST>" in reply_text:
            # Strip the tag from the message
            reply_text = reply_text.replace("<HUMAN_TAKEOVER_REQUEST>", "").strip()
            
            # Update status to confirmation
            if payload.session_id:
                await _ss.update_session_field(payload.session_id, "status", "CONFIRMING_TRANSFER")
            
            # Combine the AI's natural response with the handover prompt
            if reply_text:
                reply_text += "\n\nDo you want me to call your supervisor/HR in this chat? (yes/no)"
            else:
                reply_text = "Do you want me to call your supervisor/HR in this chat? (yes/no)"
            
            result["reply"] = reply_text

        # ── Fire-and-forget: persist to session store in background ──
        if payload.session_id:
            asyncio.create_task(_ss.save_message(payload.session_id, payload.message, result["reply"]))
            
        return ChatResponse(**result)
    except Exception as e:
        import httpx
        import traceback
        logger.error(f"Error in chat endpoint: {str(e)}\n{traceback.format_exc()}")
        
        error_msg = "⚠️ Oops! Something went wrong."
        if isinstance(e, httpx.HTTPStatusError):
            try:
                error_data = e.response.json()
                api_msg = error_data.get("error", {}).get("message", e.response.text)
            except Exception:
                api_msg = e.response.text
            error_msg = f"⚠️ API Error ({e.response.status_code}): {api_msg}"
                
        return ChatResponse(reply=error_msg, retrieved_chunks=[])


@fastapi_app.post("/api/session/update-name")
async def update_session_user_name(payload: SessionUpdateNameRequest):
    """Updates the user name and renames the session file to match the name."""
    new_id = await _ss.rename_session(payload.session_id, payload.user_name)
    return {"status": "ok", "new_session_id": new_id}


@fastapi_app.post("/api/session/init")
async def init_session(payload: SessionInitRequest):
    """
    Pre-creates a session on widget load — before the user opens the chat.
    If user_name is provided the session is immediately renamed to the canonical
    named session (merging any existing temp session).  The caller should store
    the returned session_id in localStorage so subsequent requests use the same ID.
    """
    import re

    # Ensure the temp session file exists first
    await _ss.ensure_session(payload.session_id)

    if payload.user_name and payload.user_name.strip():
        # Immediately promote to named session
        clean_name = re.sub(r'[^\w\s-]', '', payload.user_name).strip().replace(' ', '_')
        new_id = await _ss.rename_session(payload.session_id, payload.user_name)
        logger.info("session/init: pre-created named session '%s' for user '%s'", new_id, payload.user_name)
        return {"status": "ok", "session_id": new_id}
    else:
        logger.info("session/init: pre-created temp session '%s' (name unknown yet)", payload.session_id)
        return {"status": "ok", "session_id": payload.session_id}


@fastapi_app.get("/api/session/history/{session_id}")
async def get_session_history(session_id: str):
    """Returns the chat history for a session formatted for the frontend."""
    data = await _ss.get_session(session_id)
    if not data:
        return {"history": [], "last_id": 0}
    
    history = []
    messages = data.get("messages", [])
    current_chat_id = data.get("current_chat_id", 1)
    
    for msg in messages:
        if msg.get("chat_id", 1) == current_chat_id:
            if msg.get("question"):
                history.append({"role": "user", "content": msg["question"]})
            if msg.get("answer"):
                history.append({"role": "assistant", "content": msg["answer"]})
            
    max_id = max([m["id"] for m in messages if m.get("chat_id", 1) == current_chat_id]) if [m for m in messages if m.get("chat_id", 1) == current_chat_id] else 0
    return {"history": history, "user_name": data.get("user_name", ""), "status": data.get("status", "AI"), "last_id": max_id}


@fastapi_app.post("/api/chat/poll")
async def chat_poll(payload: PollRequest):
    """Frontend polls this when in HUMAN_TAKEOVER or QUEUED mode to get new admin messages."""
    sess = await _ss.get_session(payload.session_id)
    if not sess:
        return {"status": "AI", "new_messages": [], "last_id": payload.last_id}
    
    status = sess.get("status", "AI")
    messages = sess.get("messages", [])
    current_chat_id = sess.get("current_chat_id", 1)
    
    new_msgs = []
    for m in messages:
        if m.get("chat_id", 1) == current_chat_id and m["id"] > payload.last_id:
            # Only return async messages (admin replies have no 'question' field)
            if m.get("answer") and not m.get("question"):
                new_msgs.append({"id": m["id"], "role": "assistant", "content": m["answer"]})
                
    max_id = max([m["id"] for m in messages if m.get("chat_id", 1) == current_chat_id]) if [m for m in messages if m.get("chat_id", 1) == current_chat_id] else payload.last_id
    return {"status": status, "blocked": sess.get("blocked", False), "new_messages": new_msgs, "last_id": max_id}


@fastapi_app.post("/api/chat/rate")
async def chat_rate(payload: RateRequest):
    """Saves a 1-5 star rating for the session."""
    await _ss.update_session_field(payload.session_id, "rating", payload.rating)
    return {"status": "ok"}

@fastapi_app.post("/api/chat/remind")
async def chat_remind(payload: ResetRequest):
    """Sets a reminder alert for the admin panel."""
    data = await _ss.get_session(payload.session_id)
    if data:
        new_alert_id = data.get("alert_id", 0) + 1
        await _ss.update_session_field(payload.session_id, "alert_id", new_alert_id)
        await _ss.update_session_field(payload.session_id, "alert_type", "reminder")
    return {"status": "ok"}

@fastapi_app.post("/api/chat/exit_queue")
async def chat_exit_queue(payload: ResetRequest):
    """Exits the queue and goes back to AI mode."""
    data = await _ss.get_session(payload.session_id)
    if data:
        await _ss.update_session_field(payload.session_id, "status", "AI")
        await _ss.save_message(payload.session_id, "SYSTEM", "User exited the queue.")
        new_alert_id = data.get("alert_id", 0) + 1
        await _ss.update_session_field(payload.session_id, "alert_id", new_alert_id)
        await _ss.update_session_field(payload.session_id, "alert_type", "exit_queue")
    return {"status": "ok"}

@fastapi_app.post("/api/chat/reset")
async def chat_reset(payload: ResetRequest):
    """Starts a new chat thread within the same session."""
    data = await _ss.get_session(payload.session_id)
    if data:
        new_chat_id = data.get("current_chat_id", 1) + 1
        await _ss.update_session_field(payload.session_id, "current_chat_id", new_chat_id)
        await _ss.update_session_field(payload.session_id, "status", "AI")
        await _ss.update_session_field(payload.session_id, "rating", None)
    return {"status": "ok"}

@fastapi_app.post("/api/chat/delete")
async def chat_delete(payload: DeleteChatRequest):
    """Deletes the current chat thread for the user."""
    lock = await _ss._get_lock(payload.session_id)
    path = _ss._get_session_path(payload.session_id)
    async with lock:
        if path.exists():
            data = await asyncio.to_thread(_ss._read_json, path)
            current_chat_id = data.get("current_chat_id", 1)
            # Retain only messages that DO NOT belong to current_chat_id
            data["messages"] = [m for m in data.get("messages", []) if m.get("chat_id", 1) != current_chat_id]
            data["message_count"] = len(data["messages"])
            await asyncio.to_thread(_ss._write_json, path, data)
    return {"status": "ok"}



# ── Mount Gradio at /ui — MUST be last, wraps the FastAPI app ─────────────────
# All FastAPI routes (including /admin) are already registered above and are
# preserved inside the Gradio-wrapped Starlette app.
app = gr.mount_gradio_app(fastapi_app, demo, path="/ui")