import os
import json
import re
import asyncio
from contextlib import asynccontextmanager
from dotenv import load_dotenv
from operator import itemgetter

# import gradio as gr
from fastapi import FastAPI, Depends, HTTPException, Header
from fastapi.responses import JSONResponse

# Make sure you have these files in a 'utils' folder
from utils.DocsLoader import load_and_chunk
from utils.Schemas import RunRequest, RunResponse
# from concurrent.futures import ThreadPoolExecutor
from langchain_community.vectorstores import FAISS
from langchain.schema import Document
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_huggingface import HuggingFaceEmbeddings # Correct new import
# from langchain_chroma import Chroma
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
from langchain.prompts import PromptTemplate


# Load environment variables
load_dotenv()

# --- 1. Lifespan Event Handler (The New, Correct Way) ---
# This dictionary will hold our loaded models
ml_models = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    # This code runs ONCE when the application starts up
    print("🚀 Initializing models and prompt template...")

    try:
        GOOGLE_API_KEY = os.getenv("gemini_api_key")
        print("🔑 gemini_api_key:", "FOUND" if GOOGLE_API_KEY else "NOT FOUND")

        if not GOOGLE_API_KEY:
            raise RuntimeError("CRITICAL: Missing GOOGLE_API_KEY in environment secrets!")
    
        # Load models into the shared dictionary
        ml_models["embedder"] = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
        # cross_encoder_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
        cross_encoder_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-large")
        ml_models["reranker_compressor"] = CrossEncoderReranker(model=cross_encoder_model, top_n=5)
        ml_models["llm"] = ChatGoogleGenerativeAI(model="gemini-1.5-flash", api_key=GOOGLE_API_KEY)
        ml_models["prompt_template"] = PromptTemplate.from_template(
         """
        You are an expert decision maker Assistant in the domain such as insurance, legal compliance, human resources, and contract management.

The customer has submitted a query. If the query has details about age , gender ,procedure ,location , policy duration, then parse it and understand the query properly
If not, the raw question is provided instead:
- Query: {full_query}

We retrieved the following policy clauses and rules relevant to this case:
{context}

### Task:
If the details about age , gender , procedure , location , policy duration are available, do all of the following:
1. Decide whether the procedure is covered.
2. Estimate the claimable amount.
3. Justify with the relevant clause.
and answer the query  precisely as insurance agent.

Otherwise, answer the question concisely and clearly using the retrieved context.

### Output format:
If query involes age , gender , procedure , location , policy duration answer like below:
{{
  "decision": "approved / rejected",
  "amount": "INR amount or null",
  "justification": "Refer to specific clause"
  Make it a perfect and concise.
}}
Else:
{{
  "response": "Concise natural language answer"
  Make it a perfect and concise.
}}

##NOTE : Do  not mention document id or its page number just mention clauses if applicable .
        """
    )
        print("✅ Models and prompt loaded successfully!")
    except Exception as e:
        print("❌ Lifespan error:", str(e))
        raise e

    yield
    print("🧹 Cleaning up.")
    ml_models.clear()
# --- 2. FastAPI App Instance ---
# We pass the lifespan function to the FastAPI constructor
app = FastAPI(title="HackRX RAG Server", lifespan=lifespan)

# --- 3. API Key Verification ---
TEAM_API_KEY = os.getenv("TEAM_API_KEY")

def verify_api_key(authorization: str = Header(...)):
    if not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Invalid Authorization header format")
    token = authorization.split("Bearer ")[1]
    # print(token)
    # print(TEAM_API_KEY)
    if token != TEAM_API_KEY:
        raise HTTPException(status_code=403, detail="Invalid or missing API key")


# --- 4. Parsing Helper ---
def parse_llm_response(content: str) -> str:
    try:
        # Remove code fences and clean up
        content_cleaned = re.sub(r"^```json|```$", "", content.strip(), flags=re.IGNORECASE).strip()
        data = json.loads(content_cleaned)

        if isinstance(data, dict):
            if "decision" in data:
                decision = data.get("decision", "N/A").upper()
                amount = data.get("amount", "Not specified")
                justification = data.get("justification", "No justification provided.")
                return f"Decision: {decision}\nAmount: {amount}\nJustification: {justification}"

            elif "response" in data:
                return data["response"]

        return "The response was parsed but didn't match expected structure."

    except json.JSONDecodeError:
        return f"Unstructured response:\n{content.strip()}"

    except Exception as e:
        return f"An error occurred while processing the response: {str(e)}"


# --- 5. Main API Endpoint ---
@app.post("/api/v1/hackrx/run", response_model=RunResponse, dependencies=[Depends(verify_api_key)])
async def run_hackrx(req: RunRequest):
    chunks = load_and_chunk(str(req.documents))
    if not chunks:
        return JSONResponse({"error": "No documents could be processed."}, status_code=400)

        
    ####code for parallel####################################################################################

    vectorstore = await FAISS.afrom_documents(
        documents=chunks, 
        embedding=ml_models["embedder"]
    )

    dense_retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
    

    # Create retrievers using the pre-loaded models from our ml_models dictionary
    keyword_retriever = BM25Retriever.from_documents(chunks)
    keyword_retriever.k = 3
    # dense_retriever = Chroma.from_documents(documents=chunks, embedding=ml_models["embedder"]).as_retriever()
    ensemble_retriever = EnsembleRetriever(retrievers=[keyword_retriever, dense_retriever], weights=[0.4, 0.65])
    
    compression_retriever = ContextualCompressionRetriever(
        base_retriever=ensemble_retriever, base_compressor=ml_models["reranker_compressor"]
    )

    # Define the RAG chain using pre-loaded components
    hybrid_rag_chain = (
        {"context": itemgetter("full_query") | compression_retriever, "full_query": itemgetter("full_query")}
        | ml_models["prompt_template"]
        | ml_models["llm"]
    )
    
    # answers = []
    # for q in req.questions:
    #     try:
    #         result = await hybrid_rag_chain.ainvoke({"full_query": q})
    #         parsed = parse_llm_response(result.content)
    #         answers.append(parsed)
    #     except Exception as e:
    #         return JSONResponse({"error": str(e)}, status_code=500)
    tasks = [hybrid_rag_chain.ainvoke({"full_query": q}) for q in req.questions]
    results = await asyncio.gather(*tasks)
    
    # Extract the content from each result and parse it
    answers = [parse_llm_response(result.content) for result in results]

    return JSONResponse({"answers": answers}, status_code=200)

@app.get("/", include_in_schema=False)
def root():
    return {"message": "API is running. Go to /docs for documentation."}


# def dummy_gradio(): return "✅ API running!"
# gr.Interface(fn=dummy_gradio, inputs=[], outputs="text").launch(server_name="0.0.0.0", server_port=7860)