"""
app.py — Vibhu Solutions AI Agent — FastAPI Backend

Start  : python app.py
Docs   : http://localhost:8500/docs
Health : http://localhost:8500/health
"""

import os
import uvicorn
from contextlib import asynccontextmanager
from pathlib import Path
from dotenv import load_dotenv
from typing import List, Dict, Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel

import firebase_admin
from firebase_admin import credentials, firestore

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate

load_dotenv(override=True)

# Firebase Initialization
_db = None
def _get_db():
    global _db
    if _db is None:
        if not firebase_admin._apps:
            fb_creds = os.getenv("FIREBASE_CREDENTIALS", "").strip()
            if fb_creds:
                try:
                    import json
                    cred = credentials.Certificate(json.loads(fb_creds))
                    firebase_admin.initialize_app(cred)
                    _db = firestore.client()
                    print("✅  Firebase connected.")
                except Exception as e:
                    print(f"⚠️  Firebase init failed: {e} — leads stored in memory only.")
            else:
                print("⚠️  FIREBASE_CREDENTIALS not set — leads stored in memory only.")
    return _db

# Configuration 
GROQ_KEYS = [k.strip() for k in [
    os.getenv("GROQ_API_KEY",   ""),
    os.getenv("GROQ_API_KEY_2", ""),
    os.getenv("GROQ_API_KEY_3", ""),
] if k.strip()]
BUSINESS_NAME = os.getenv("BUSINESS_NAME", "Vibhu Solutions")
LLM_MODEL     = os.getenv("LLM_MODEL", "llama-3.1-8b-instant")
FAISS_DIR     = "./faiss_db"
EMBED_MODEL   = "sentence-transformers/all-MiniLM-L6-v2"
MAX_HISTORY   = 8    # last 8 messages (~4 exchanges) per session
MAX_SESSIONS  = 500  # evict oldest sessions beyond this limit

# In-memory session store  {session_id: [{"role": "user"|"assistant", "content": "..."}]}
session_memory: Dict[str, List[Dict[str, str]]] = {}
_active_key_i: int = 0

# In-memory leads store  [{name, email, phone, time}]
leads_store: List[Dict[str, str]] = []

# ── Agent System Prompt (system message only — question goes in human message)
SYSTEM_PROMPT = """\
You are an intelligent AI Agent for {business_name}, a professional technology company \
based in Bengaluru, India specialising in AI solutions, software development, cloud, \
blockchain, IoT, mobile apps, DevOps, UI/UX, and digital marketing.

RULES — FOLLOW STRICTLY:

RULE 0 — GREETINGS:
If the user sends a greeting (hi, hello, hey, how are you, good morning, etc.),
respond warmly and briefly, then invite them to ask about {business_name}'s services.
Do NOT treat greetings as out-of-scope questions.

RULE 1 — SCOPE:
Only discuss topics related to {business_name}: services, technologies, team, industries,
process, timelines, support, and contact information.
For unrelated topics (general knowledge, news, coding tutorials, math, etc.) politely decline.
Vary the wording each time — do not use the same phrase repeatedly.
After two redirects on the same topic, apply RULE 5.

RULE 2 — PRICING (NEVER REVEAL):
Never mention prices, cost estimates, or budget figures. Always respond:
  "For accurate pricing, please contact our team:

   📞 +91 9380345108
   📧 contact@vibhusolutions.com
   📍 #57, 2nd floor, 2nd cross, 80 Feet Road, Bhuvaneshwari Nagar,
      5th Block, BSK 3rd Stage, Bengaluru – 560070

   Our team will provide a custom quote based on your requirements."

RULE 3 — ACCURACY:
Use ONLY information from the KNOWLEDGE BASE. Do not invent project names or facts.
If not in the knowledge base: "I don't have that detail. Contact: contact@vibhusolutions.com | +91 9380345108"

RULE 4 — LEAD CAPTURE:
When a user wants a quote or consultation:
  A) If name/email/phone already in USER CONTEXT, confirm them — do NOT ask again.
  B) If missing, ask for only the missing info, one at a time.
  C) Once confirmed: "Thank you! Our team will contact you within a few hours."

RULE 5 — HUMAN HANDOFF:
If user wants a human, is frustrated, or has complex contract/legal questions:
  "Let me connect you with our team:

   📞 +91 9380345108
   📧 contact@vibhusolutions.com
   📍 #57, 2nd floor, 2nd cross, 80 Feet Road, Bhuvaneshwari Nagar,
      5th Block, BSK 3rd Stage, Bengaluru – 560070"

RULE 6 — CONTACT INFO FORMAT:
When a user asks for contact details, address, location, or "how to reach you", always reply with this exact block:

   📞 Phone / WhatsApp : +91 9380345108
   📧 Email            : contact@vibhusolutions.com
   📍 Address          : #57, 2nd floor, 2nd cross, 80 Feet Road,
                         Bhuvaneshwari Nagar, 5th Block, BSK 3rd Stage,
                         Bengaluru – 560070
   🌐 Website          : vibhusolutions.com

   We respond within a few hours. For the fastest reply, WhatsApp us!

RULE 7 — TONE:
Be warm, professional, concise. Use bullet points for readability.
Do NOT start every reply with the user's name. Just answer directly.
Only include the full contact block when genuinely relevant (contact/location questions, pricing, human handoff).

---
CONVERSATION HISTORY (context only — do not repeat):
{history}

KNOWLEDGE BASE:
{context}

USER CONTEXT:
{user_context}
"""

#  Global RAG components (loaded once at startup)
_retriever = None
_llm       = None
_prompt    = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Lazy load: don't block startup — components load on first chat request
    print(f"✅  {BUSINESS_NAME} AI Agent server started. Components load on first request.")
    yield


def _make_llm() -> "ChatGroq":
    return ChatGroq(
        groq_api_key=GROQ_KEYS[_active_key_i % len(GROQ_KEYS)],
        model_name=LLM_MODEL,
        temperature=0.2,
    )


def _build_components():
    if not GROQ_KEYS:
        raise ValueError("GROQ_API_KEY not set. Add it to your .env file.")
    if not Path(FAISS_DIR).exists():
        raise FileNotFoundError(
            "FAISS index not found. Run 'python ingest.py' first."
        )

    embeddings = HuggingFaceEmbeddings(
        model_name=EMBED_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )
    vectorstore = FAISS.load_local(
        FAISS_DIR, embeddings, allow_dangerous_deserialization=True
    )
    retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
    llm       = _make_llm()
    prompt    = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        ("human", "{question}"),
    ])
    return retriever, llm, prompt


def _format_history(messages: List[Dict[str, str]]) -> str:
    if not messages:
        return "None"
    lines = []
    for m in messages:
        role = "user" if m["role"] == "user" else "assistant"
        content = m["content"]
        if role == "assistant" and len(content) > 300:
            content = content[:300] + "…"
        lines.append(f"{role}: {content}")
    return "\n".join(lines)


#  FastAPI app 
app = FastAPI(
    title=f"{BUSINESS_NAME} AI Agent",
    version="2.0.0",
    lifespan=lifespan,
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

static_dir = Path("static")
static_dir.mkdir(exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")


#  Pydantic models 
class ChatRequest(BaseModel):
    message:    str
    session_id: str = "default"
    user_name:  Optional[str] = ""
    user_email: Optional[str] = ""
    user_phone: Optional[str] = ""


class ChatResponse(BaseModel):
    reply:      str
    session_id: str


#  Endpoints 
@app.get("/")
def root():
    page = Path("static/index.html")
    if page.exists():
        return FileResponse(page)
    return {
        "status":  "running",
        "message": f"{BUSINESS_NAME} AI Agent is live!",
        "docs":    "/docs",
        "chat":    "POST /chat",
    }


@app.get("/health")
def health():
    return {
        "status":      "ok",
        "agent_ready": _retriever is not None,
        "business":    BUSINESS_NAME,
        "model":       LLM_MODEL,
    }


@app.get("/leads")
def get_leads(key: str = ""):
    LEADS_KEY = os.getenv("LEADS_API_KEY", "vibhu2025")
    if key != LEADS_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key.")
    db = _get_db()
    if db:
        docs = db.collection("leads").order_by("time", direction=firestore.Query.DESCENDING).stream()
        leads = [doc.to_dict() for doc in docs]
    else:
        leads = leads_store
    return {"total": len(leads), "leads": leads}


@app.post("/chat", response_model=ChatResponse)
def chat(request: ChatRequest):
    global _retriever, _llm, _prompt, _active_key_i
    if not request.message.strip():
        raise HTTPException(status_code=400, detail="Message cannot be empty.")

    # Lazy-load components on first request
    if _retriever is None or _llm is None:
        try:
            _retriever, _llm, _prompt = _build_components()
            print(f"✅  {BUSINESS_NAME} AI Agent is ready (loaded on first request).")
        except (ValueError, FileNotFoundError):
            return ChatResponse(
                reply=(
                    "The AI Agent is not configured yet. "
                    "Please set GROQ_API_KEY and run 'python ingest.py'."
                ),
                session_id=request.session_id,
            )

    # Capture lead if name+email provided and not already stored
    if request.user_name and request.user_email:
        from datetime import datetime
        lead_data = {
            "name":  request.user_name,
            "email": request.user_email,
            "phone": request.user_phone or "",
            "time":  datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC"),
        }
        db = _get_db()
        if db:
            doc_ref = db.collection("leads").document(request.user_email)
            if not doc_ref.get().exists:
                doc_ref.set(lead_data)
        else:
            already = any(l["email"] == request.user_email for l in leads_store)
            if not already:
                leads_store.append(lead_data)

    try:
        # Retrieve session history
        history = session_memory.get(request.session_id, [])

        # Retrieve relevant knowledge-base chunks
        docs    = _retriever.invoke(request.message.strip())
        context = "\n\n".join(doc.page_content for doc in docs)

        # Build user-context block (used by RULE 4 to skip re-asking known details)
        ctx_parts = []
        if request.user_name:  ctx_parts.append(f"Name  : {request.user_name}")
        if request.user_email: ctx_parts.append(f"Email : {request.user_email}")
        if request.user_phone: ctx_parts.append(f"Phone : {request.user_phone}")
        user_ctx = "\n".join(ctx_parts) if ctx_parts else "Not provided"

        # Format prompt
        formatted = _prompt.format_messages(
            business_name=BUSINESS_NAME,
            context=context,
            question=request.message.strip(),
            history=_format_history(history),
            user_context=user_ctx,
        )

        # Invoke LLM — auto-rotate key on rate limit (if multiple keys available)
        try:
            reply = _llm.invoke(formatted).content
        except Exception as llm_err:
            err_str = str(llm_err)
            if ("rate_limit_exceeded" in err_str or "429" in err_str) and len(GROQ_KEYS) > 1:
                _active_key_i = (_active_key_i + 1) % len(GROQ_KEYS)
                _llm = _make_llm()
                print(f"⚠️  Rate limit hit — rotated to Groq key #{_active_key_i + 1}")
                reply = _llm.invoke(formatted).content
            else:
                raise

        # Persist history (capped at MAX_HISTORY messages)
        history.append({"role": "user",      "content": request.message.strip()})
        history.append({"role": "assistant", "content": reply})
        session_memory[request.session_id] = history[-MAX_HISTORY:]

        # Evict oldest sessions if memory grows too large
        if len(session_memory) > MAX_SESSIONS:
            oldest = next(iter(session_memory))
            del session_memory[oldest]

        return ChatResponse(reply=reply, session_id=request.session_id)

    except Exception as e:
        err = str(e)
        # Rate limit
        if "rate_limit_exceeded" in err or "429" in err:
            print(f"⚠️  Rate limit: {err}")
            return ChatResponse(
                reply=(
                    "Our AI assistant is temporarily busy. "
                    "Please try again in a moment.\n\n"
                    "Or contact us directly:\n"
                    "📧 contact@vibhusolutions.com\n"
                    "📞 +91 9380345108"
                ),
                session_id=request.session_id,
            )
        # Network / connection error
        if "ConnectError" in err or "getaddrinfo" in err or "APIConnectionError" in err or "Connection error" in err:
            print(f"⚠️  Network error (transient): {err}")
            return ChatResponse(
                reply=(
                    "I'm having trouble connecting right now. "
                    "Please try again in a moment.\n\n"
                    "Or reach us directly:\n"
                    "📧 contact@vibhusolutions.com\n"
                    "📞 +91 9380345108"
                ),
                session_id=request.session_id,
            )
        # Unexpected error — log full traceback
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Agent error: {err}")


if __name__ == "__main__":
    import asyncio
    from contextlib import nullcontext
    port = int(os.getenv("PORT", 8500))   # 7860 on HuggingFace Spaces, 8000 locally
    config = uvicorn.Config("app:app", host="0.0.0.0", port=port, reload=False)
    server = uvicorn.Server(config)
    server.capture_signals = nullcontext   # Fix Python 3.14 signal handling crash
    asyncio.run(server.serve())