import os
import re
import time
import json
import difflib
from typing import List, Optional, Dict, Any
from contextlib import asynccontextmanager

import numpy as np
import pandas as pd
import faiss
from sentence_transformers import SentenceTransformer, CrossEncoder
from groq import Groq

from fastapi import FastAPI, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

# ==========================================
# Paths & Default Configurations
# ==========================================
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
FAISS_INDEX_PATH = os.path.join(BASE_DIR, "questions.index")
DATASET_PATH = os.path.join(BASE_DIR, "notebooks", "AHD_english_cleaned.xlsx")
SYNONYMS_PATH = os.path.join(BASE_DIR, "rag_model", "models", "medical_synonyms.json")
CONFIG_PATH = os.path.join(BASE_DIR, "rag_model", "rag_config.json")

DEFAULT_GROQ_API_KEY = os.getenv("GROQ_API_KEY", "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")
LLM_MODEL = os.getenv("LLM_MODEL", "openai/gpt-oss-20b")

# Load config if available
if os.path.exists(CONFIG_PATH):
    with open(CONFIG_PATH, "r", encoding="utf-8") as f:
        RAG_CONFIG = json.load(f)
else:
    RAG_CONFIG = {
        "embedding_model": "all-MiniLM-L6-v2",
        "reranker_model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
        "retrieve_top_n": 20,
        "rerank_top_k": 6,
    }

# Dynamic Global State
rag_resources: Dict[str, Any] = {}

SYSTEM_PROMPT = """
You are a medical RAG assistant specialized ONLY in:
- Diabetes
- Endocrinology
- Thyroid disorders
- Hormonal disorders
- Endocrine glands and related disorders
- Nutrition, glucose management, insulin, and complications DIRECTLY related to diabetes/endocrinology

You operate in STRICT RAG MODE.

==================================================
CRITICAL DOMAIN RULE
==================================================
The USER QUESTION itself determines whether the question is in-domain.
A question is IN-DOMAIN only if its actual subject is diabetes, endocrinology,
thyroid, hormones, endocrine glands/disorders, or a problem explicitly stated
by the user to be related to diabetes/endocrinology.

If the user's actual question is about another body system or another medical
specialty, it is OUT-OF-DOMAIN.

==================================================
TYPE 1 — IN-DOMAIN MEDICAL QUESTION
==================================================
If and ONLY IF the USER QUESTION itself is in-domain:
- Answer using ONLY the RETRIEVED MEDICAL EVIDENCE.
- Do NOT use pretrained/background medical knowledge.
- Do NOT guess or invent missing details.

==================================================
DOSAGE & TREATMENT SAFETY — MANDATORY RULES
==================================================
These rules OVERRIDE everything else for any answer involving medications,
insulin, doses, units, quantities, or treatment amounts:

1. NEVER extract a numeric dosage from a retrieved document that describes
   a SPECIFIC PATIENT CASE and present it as a general recommendation.

2. A dosage found in evidence is ONLY valid for the exact clinical scenario
   described in that document (e.g., specific blood glucose level, patient
   weight, insulin type, diabetes type, age, medical history).

3. If the user's question is GENERAL (e.g., "How much insulin should I take?")
   and the evidence only contains case-specific dosages, you MUST NOT quote
   those numbers as a general answer.

4. For ANY medication dosage question where the evidence is:
   - Case-specific (belongs to a specific patient scenario), OR
   - Contradicted by other evidence, OR
   - Insufficient to determine a safe dose for the user's exact situation:
   → Explicitly state that the dose CANNOT be determined from the available
     information without individual physician assessment.
   → Always direct the user to consult their treating physician.

5. Do NOT present partial evidence (one document saying "4-8 units") as a
   complete answer when other evidence clearly states that doses are
   patient-specific and require physician supervision.

EXAMPLE — WRONG (do not do this):
User: "How much insulin should I take?"
Evidence has: "If sugar > 300, inject 4-8 units"
WRONG answer: "You may need 4-8 units of insulin."

EXAMPLE — CORRECT:
User: "How much insulin should I take?"
CORRECT answer: "Insulin doses cannot be determined from general guidelines.
They depend on your blood sugar level, weight, type of diabetes, insulin type,
and medical history. You must consult your treating physician to determine the
appropriate dose for your specific situation."

==================================================
TYPE 2 — APP / IDENTITY QUESTIONS
==================================================
Only for direct questions about the assistant/application itself:
Answer briefly and naturally.

==================================================
TYPE 3 — EVERYTHING ELSE / OUT OF DOMAIN / INSUFFICIENT EVIDENCE
==================================================
For ALL TYPE 3 cases, reply with ONLY one short refusal:
If the user wrote in Arabic:
"معرفش، السؤال ده مش جزء من تخصصي (السكر والغدد الصماء)."
If the user wrote in English:
"I don't know — this is outside my specialty (diabetes and endocrinology)."

==================================================
FINAL SAFETY NOTE FOR TYPE 1 ONLY
==================================================
At the end of every TYPE 1 medical answer, write:
"This information is based on the available medical evidence and is for general informational purposes. It is not a diagnosis or a substitute for professional medical advice."
"""

# ==========================================
# RAG Helper Logic
# ==========================================
def expand_query(query: str, synonyms: dict) -> str:
    lower_q = query.lower()
    extra_terms = [
        med_term
        for phrase, med_term in synonyms.items()
        if phrase in lower_q and med_term not in lower_q
    ]
    if not extra_terms:
        return query
    return f"{query} ({', '.join(dict.fromkeys(extra_terms))})"

def vector_search(query: str, top_n: int = 20):
    embedder = rag_resources["embedder"]
    faiss_index = rag_resources["faiss_index"]
    metadata_store = rag_resources["metadata_store"]
    synonyms = rag_resources["synonyms"]

    query_for_embedding = expand_query(query, synonyms)
    q_emb = embedder.encode(
        [query_for_embedding],
        convert_to_numpy=True,
        normalize_embeddings=True
    ).astype("float32")

    scores, idxs = faiss_index.search(q_emb, top_n)
    retrieved = []
    for score, idx in zip(scores[0], idxs[0]):
        if idx == -1:
            continue
        payload = metadata_store[int(idx)]
        retrieved.append({
            "doc_id": payload["doc_id"],
            "Question": payload["question"],
            "Answer": payload["answer"],
            "Category": payload["category"],
            "similarity": float(score),
        })
    return pd.DataFrame(retrieved)

def rerank(query: str, candidates: pd.DataFrame, top_k: int = 6, alpha: float = 0.6) -> pd.DataFrame:
    if candidates.empty:
        return candidates

    reranker = rag_resources["reranker"]
    pairs = [
        (query, f"Question: {row['Question']}\nAnswer: {row['Answer']}")
        for _, row in candidates.iterrows()
    ]

    rerank_scores = reranker.predict(pairs)
    reranked = candidates.copy()
    reranked["rerank_score"] = rerank_scores

    def norm(s):
        s = s.astype(float)
        rng = s.max() - s.min()
        return (s - s.min()) / rng if rng > 0 else s * 0

    reranked["sim_norm"] = norm(reranked["similarity"])
    reranked["rerank_norm"] = norm(reranked["rerank_score"])
    reranked["final_score"] = alpha * reranked["rerank_norm"] + (1 - alpha) * reranked["sim_norm"]

    return reranked.sort_values("final_score", ascending=False).head(top_k).reset_index(drop=True)

def is_near_duplicate(text_a: str, text_b: str, threshold: float = 0.92) -> bool:
    return difflib.SequenceMatcher(None, text_a, text_b).ratio() > threshold

def build_evidence(reranked: pd.DataFrame, max_answer_chars: int = 700) -> List[dict]:
    evidence = []
    seen_answers = []

    for _, row in reranked.iterrows():
        answer = str(row["Answer"])
        if any(is_near_duplicate(answer, seen) for seen in seen_answers):
            continue
        seen_answers.append(answer)

        evidence.append({
            "doc_id": int(row["doc_id"]),
            "question": str(row["Question"]),
            "answer": answer[:max_answer_chars] + ("..." if len(answer) > max_answer_chars else ""),
            "category": str(row["Category"]),
            "similarity": round(float(row["similarity"]), 3),
            "rerank_score": round(float(row["rerank_score"]), 3),
        })
    return evidence

def generate_answer(query: str, user_data: str, evidence: List[dict]) -> str:
    groq_client = rag_resources["groq_client"]
    if not evidence:
        return "I don't know — this is outside my specialty (diabetes and endocrinology)."

    evidence_blocks = [
        f"[{e['doc_id']}] (category: {e['category']})\nRelated question: {e['question']}\nAnswer: {e['answer']}"
        for e in evidence
    ]
    evidence_text = "\n\n".join(evidence_blocks)

    user_message = f"""
USER QUESTION:
{query}

USER DATA:
{user_data}

RETRIEVED MEDICAL EVIDENCE:
{evidence_text}

TASK:
Answer the USER QUESTION using ONLY the RETRIEVED MEDICAL EVIDENCE.
Cite every important medical claim using the relevant [doc_id].
If the evidence is insufficient or irrelevant, give a clear refusal.
"""

    try:
        response = groq_client.chat.completions.create(
            model=LLM_MODEL,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_message}
            ],
            temperature=0.1,
            max_tokens=800,
        )
        answer = response.choices[0].message.content
        answer = re.sub(r'\[\d+\]', '', answer)
        return answer   
    except Exception as e:
        # Fallback or error handling for Groq API
        return f"Error generating answer from LLM: {str(e)}"

# ==========================================
# FastAPI Lifecycle & App Initialization
# ==========================================
@asynccontextmanager
async def lifespan(app: FastAPI):
    print("Loading RAG resources...")

    # Load FAISS index
    if not os.path.exists(FAISS_INDEX_PATH):
        raise RuntimeError(f"FAISS index file not found at: {FAISS_INDEX_PATH}")
    faiss_index = faiss.read_index(FAISS_INDEX_PATH)

    # Load Dataset for Metadata
    if not os.path.exists(DATASET_PATH):
        raise RuntimeError(f"Dataset file not found at: {DATASET_PATH}")
    df = pd.read_excel(DATASET_PATH, engine="openpyxl").reset_index(drop=True)
    df["doc_id"] = df.index
    metadata_store = [
        {
            "doc_id": row["doc_id"],
            "question": row["Question"],
            "answer": row["Answer"],
            "category": row["Category"],
        }
        for _, row in df.iterrows()
    ]

    # Load Synonyms
    synonyms = {}
    if os.path.exists(SYNONYMS_PATH):
        with open(SYNONYMS_PATH, "r", encoding="utf-8") as f:
            synonyms = json.load(f)

    # Load Transformer Models
    embedder = SentenceTransformer(RAG_CONFIG.get("embedding_model", "all-MiniLM-L6-v2"))
    reranker = CrossEncoder(RAG_CONFIG.get("reranker_model", "cross-encoder/ms-marco-MiniLM-L-6-v2"))

    # Load Groq Client
    groq_client = Groq(api_key=DEFAULT_GROQ_API_KEY)

    # Save to global state
    rag_resources["faiss_index"] = faiss_index
    rag_resources["metadata_store"] = metadata_store
    rag_resources["synonyms"] = synonyms
    rag_resources["embedder"] = embedder
    rag_resources["reranker"] = reranker
    rag_resources["groq_client"] = groq_client

    print(f"RAG resources loaded successfully! Index size: {faiss_index.ntotal}")
    yield
    print("Shutting down RAG resources.")
    rag_resources.clear()

app = FastAPI(
    title="CortexRAG Medical API",
    description="Production-ready FastAPI interface for CortexRAG Medical Question Answering Engine.",
    version="1.0.0",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ==========================================
# Request / Response Schemas
# ==========================================
class QueryRequest(BaseModel):
    question: str = Field(..., example="What are the common symptoms of hypothyroidism?", min_length=2)
    user_data: Optional[str] = Field(default="", example="Age: 45, Gender: Female")
    top_k: Optional[int] = Field(default=6, ge=1, le=20)

class EvidenceItem(BaseModel):
    question: str
    answer: str
    category: str

class QueryResponse(BaseModel):
    status: str
    question: str
    answer: str
    execution_time_sec: float

# ==========================================
# Endpoints
# ==========================================
@app.get("/", tags=["Health"])
def root():
    return {
        "status": "online",
        "service": "CortexRAG Medical API",
        "documentation": "/docs"
    }

@app.get("/health", tags=["Health"])
def health_check():
    return {
        "status": "healthy",
        "faiss_index_entries": rag_resources.get("faiss_index", None).ntotal if "faiss_index" in rag_resources else 0,
        "model_loaded": "embedder" in rag_resources
    }

def _run_rag(payload: QueryRequest) -> QueryResponse:
    start_time = time.time()
    try:
        candidates = vector_search(payload.question, top_n=RAG_CONFIG.get("retrieve_top_n", 20))
        reranked_df = rerank(payload.question, candidates, top_k=payload.top_k)
        evidence = build_evidence(reranked_df)
        answer = generate_answer(payload.question, payload.user_data or "", evidence)
        elapsed = round(time.time() - start_time, 3)
        return QueryResponse(
            status="success",
            question=payload.question,
            answer=answer,
            execution_time_sec=elapsed
        )
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"An error occurred during inference: {str(e)}"
        )

@app.post("/query", response_model=QueryResponse, tags=["RAG Inference"])
def query(payload: QueryRequest):
    return _run_rag(payload)

@app.post("/predict", response_model=QueryResponse, tags=["RAG Inference"])
def predict(payload: QueryRequest):
    return _run_rag(payload)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="127.0.0.1", port=8000, reload=True)