from fastapi import FastAPI, Request, HTTPException

from models.schemas import ChatRequest

from cache.user_manager import UserManager
from cache.response_cache import ResponseCache

from admin.routes import router as admin_router

from rag.engine import RAGEngine


# ==========================================================
# FASTAPI
# ==========================================================

app = FastAPI(
    title="AIForge",
    version="1.1.0"
)


# ==========================================================
# INITIALIZE SERVICES (Load Once)
# ==========================================================

engine = RAGEngine()

user_manager = UserManager()

response_cache = ResponseCache()


# ==========================================================
# ROUTERS
# ==========================================================

app.include_router(admin_router)


# ==========================================================
# HOME
# ==========================================================

@app.get("/")
def home():

    return {
        "message": "🚀 AIForge Running Successfully"
    }


# ==========================================================
# CHAT
# ==========================================================

@app.post("/chat")
def chat(
    data: ChatRequest,
    request: Request
):

    prompt = data.prompt.strip()

    # -------------------------------------------------------
    # USER IP
    # -------------------------------------------------------

    ip = request.headers.get(
        "X-Forwarded-For",
        request.client.host
    ).split(",")[0].strip()

    # -------------------------------------------------------
    # REGISTER USER
    # -------------------------------------------------------

    user_manager.register_request(ip)

    # -------------------------------------------------------
    # CACHE CHECK
    # -------------------------------------------------------

    cached = response_cache.get(
        ip,
        prompt
    )

    if cached:

        user = user_manager.get_user(ip)

        return {

            "success": True,

            "cached": True,

            "provider": "cache",

            "remaining_questions": user["remaining"],

            "answer": cached["answer"]

        }

    # -------------------------------------------------------
    # CACHE MISS
    # -------------------------------------------------------

    # -------------------------------------------------------
    # RATE LIMIT
    # -------------------------------------------------------

    allowed, user = user_manager.allow_request(ip)

    if not allowed:

        raise HTTPException(

            status_code=429,

            detail=f"Daily limit reached. ({user['llm_requests']}/10)"

        )

    # -------------------------------------------------------
    # RAG ENGINE
    # -------------------------------------------------------

    answer = engine.chat(prompt)

    # -------------------------------------------------------
    # SAVE CACHE
    # -------------------------------------------------------

    response_cache.save(
        prompt,
        answer
    )

    # -------------------------------------------------------
    # RESPONSE
    # -------------------------------------------------------

    return {

        "success": True,

        "cached": False,

        "provider": "rag",

        "remaining_questions": user["remaining"],

        "answer": answer

    }