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

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 
# from langchain_chroma import Chroma
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker 
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
from langchain.prompts import ChatPromptTemplate

import os
from sentence_transformers import SentenceTransformer


MODEL_DIR = os.path.join("/tmp", "e5-large-v2")

if not os.path.exists(MODEL_DIR):
    print("📦 Downloading SentenceTransformer model...")
    model = SentenceTransformer("intfloat/e5-large-v2")
    model.save(MODEL_DIR)
    print("✅ Model saved at", MODEL_DIR)

    
# Load environment variables
load_dotenv()


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", #better but lil slower
                                                      # model_name="intfloat/e5-large-v2", #lil faster  but  dont know response is slow 
                                                      encode_kwargs={
                                                          "batch_size": 64,
                                                          # "normalize_embeddings": True
                                                      }
                                                    
    )
        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=6)
        ml_models["llm"] = ChatGoogleGenerativeAI(
                               # model="gemini-1.5-pro",
                               model="gemini-2.0-flash",
                               api_key=GOOGLE_API_KEY,
                               temperature=0.1,
                               max_output_tokens=300
                           )
        ml_models["prompt_template"] = ChatPromptTemplate.from_template("""
**Role**: You are an expert decision maker Assistant in the domain such as insurance, legal compliance, human resources, and contract management using ONLY the provided context.
Do not use your own knowledge.
**Context**: 
{context}

**Query**: {full_query}

**Response**:

Example how to answer for thr query:
1. query:What is the grace period for premium payment under the National Parivar Mediclaim Plus Policy?
   response: A grace period of thirty days is provided for premium payment after the due date to renew or continue the policy without losing continuity benefits.

2. query: What is the waiting period for pre-existing diseases (PED) to be covered?
   response: There is a waiting period of thirty-six (36) months of continuous coverage from the first policy inception for pre-existing diseases and their direct complications to be covered.

3. query: Are the medical expenses for an organ donor covered under this policy?
   response : Yes, the policy indemnifies the medical expenses for the organ donor's hospitalization for the purpose of harvesting the organ, provided the organ is for an insured person and the donation complies with the Transplantation of Human Organs Act, 1994.

4. query: Does this policy cover maternity expenses, and what are the conditions?
   response: Yes, the policy covers maternity expenses, including childbirth and lawful medical termination of pregnancy. To be eligible, the female insured person must have been continuously covered for at least 24 months. The benefit is limited to two deliveries or terminations during the policy period.
"""
)
        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]
    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)

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

    # dense_retriever = vectorstore.as_retriever(search_type="mmr",search_kwargs={"k": 8})
    dense_retriever = vectorstore.as_retriever(search_type="mmr",search_kwargs={"k": 8 ,"lambda_mult": 0.5})
    

    # Create retrievers using the pre-loaded models from our ml_models dictionary
    keyword_retriever = BM25Retriever.from_documents(chunks)
    keyword_retriever.k = 5
    # dense_retriever = Chroma.from_documents(documents=chunks, embedding=ml_models["embedder"]).as_retriever()
    ensemble_retriever = EnsembleRetriever(retrievers=[keyword_retriever, dense_retriever], weights=[0.35, 0.65])
    ### to make it faster we are now using our built reranker thats why commenting the code below
    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"]
    )

    tasks = [hybrid_rag_chain.ainvoke({"full_query": q}) for q in req.questions]
    results = await asyncio.gather(*tasks)
    # answers = [parse_llm_response(result.content) for result in results]
    answers = []

    for msg in results:
        # Safely access the content field
        if hasattr(msg, "content"):
            answers.append(msg.content.strip())    

    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."}