import warnings
warnings.filterwarnings("ignore", category=FutureWarning)

import os
import logging
from contextlib import asynccontextmanager
from typing import List, Dict
from dotenv import load_dotenv
from fastapi import FastAPI, Request, Depends, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel, HttpUrl

load_dotenv()
# --- Configuration ---
# Set up basic logging to see output in Hugging Face logs
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load the token directly from environment variables (set in HF Secrets)
# This is the correct way for server deployments.
EXPECTED_TOKEN = os.getenv("BEARER_TOKEN")
if not EXPECTED_TOKEN:
    logger.warning("BEARER_TOKEN environment variable not set. Using a default for local testing.")
    EXPECTED_TOKEN = "default-secret-for-testing"

http_bearer_scheme = HTTPBearer()

# This dictionary will act as a cache for our heavy models.
# It will be populated during the application's startup.
model_cache: Dict[str, any] = {}

# --- Lifespan Management (The Core Fix) ---
# This function runs ONCE when the application starts up.
# It loads the models before the API starts accepting requests.
@asynccontextmanager
async def lifespan(app: FastAPI):
    # --- Startup ---
    logger.info("--- Application Startup: Loading models... ---")
    
    # IMPORTANT: Import your agent and model-loading modules INSIDE the lifespan
    # to prevent them from being loaded at the top level.
    from agent import agent_app, set_retriever_for_tool
    
    # Store the initialized agent app in our cache so endpoints can access it.
    model_cache["agent_app"] = agent_app
    model_cache["set_retriever_for_tool"] = set_retriever_for_tool
    
    logger.info("--- Models Loaded. Application is ready to serve requests. ---")
    
    yield  # The application is now running and serving requests
    
    # --- Shutdown ---
    logger.info("--- Application Shutdown: Clearing model cache. ---")
    model_cache.clear()


# --- FastAPI App Initialization ---
# We pass our lifespan manager to the FastAPI app.
app = FastAPI(
    title="Intelligent Query-Retrieval System",
    description="An API for processing documents and answering questions using a LangGraph ReAct agent.",
    version="1.0.0",
    lifespan=lifespan
)


# --- Authentication Dependency ---
async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(http_bearer_scheme)):
    if not credentials or credentials.scheme != "Bearer":
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid authentication scheme.")
    if credentials.credentials != EXPECTED_TOKEN:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token.")
    return credentials.credentials


# --- Pydantic Models ---
class QueryRequest(BaseModel):
    documents: HttpUrl
    questions: List[str]

class QueryResponse(BaseModel):
    answers: List[str]


# --- API Endpoints ---
@app.api_route("/", methods=["GET", "HEAD"])
def read_root():
    return {"message": "API is running. Use /hackrx/run for document processing."}

@app.post("/hackrx/run", response_model=QueryResponse, tags=["Query"])
async def run_submission(request: QueryRequest, token: str = Depends(verify_token)):
    try:
        # Retrieve the pre-loaded agent from our cache
        agent_app = model_cache.get("agent_app")
        set_retriever_for_tool = model_cache.get("set_retriever_for_tool")

        if not agent_app or not set_retriever_for_tool:
            raise HTTPException(status_code=503, detail="Models are not available. The application might be starting up.")

        logger.info(f"Processing document from URL: {request.documents}")
        # Process the document ONCE per request
        from utils import process_documents
        vector_store = process_documents(str(request.documents))
        retriever = vector_store.as_retriever(search_kwargs={"k": 5})
        
        # Configure the agent's tool with the newly created retriever for this specific request
        set_retriever_for_tool(retriever)
        
        responses = []
        logger.info(f"Answering {len(request.questions)} questions...")
        for i, question in enumerate(request.questions):
            logger.info(f"Invoking agent for question {i+1}/{len(request.questions)}: '{question}'")
            
            final_state = agent_app.invoke({
                "question": question,
                "messages": [("user", question)]
            })
            
            answer_str = final_state.get("answer", "Error: Could not generate an answer.")
            responses.append(answer_str)
            logger.info(f"Generated answer for question {i+1}")

        logger.info("All questions answered successfully.")
        return QueryResponse(answers=responses)
        
    except Exception as e:
        logger.error(f"An unhandled error occurred in /hackrx/run: {e}", exc_info=True)
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))