import os

# Belt-and-suspenders against any accidental Hugging Face Hub network calls -
# every model below is loaded from a local path, this just makes sure
# transformers/sentence-transformers never try to phone home.
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

import re
import pickle
import datetime
import numpy as np
import faiss
import torch

from threading import Thread
from sentence_transformers import SentenceTransformer, CrossEncoder
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

try:
    from dateutil import parser as _dateutil_parser
    _HAS_DATEUTIL = True
except ImportError:
    _HAS_DATEUTIL = False


# ==================================================
# CONFIG
# ==================================================

INDEX_PATH = "hbl_vector_store/hbl_faiss.index"
META_PATH = "hbl_vector_store/hbl_metadata.pkl"
EMBED_MODEL = "./bge-m3"
RERANKER_MODEL_PATH = "./bge-reranker-v2-m3"   # local only, no HF downloads
LLM_MODEL_PATH = "./qwen2.5-0.5b-instruct"     # point this at your local download folder
MIN_SCORE = 0.58        # bi-encoder recall gate (coarse - the cross-encoder does the real filtering)
HISTORY_TURNS = 3       # number of past user/assistant exchanges kept for classifier context
MAX_NEW_TOKENS = 512

# --- retrieval pipeline config -----------------------------------------
# Query -> BGE-M3 embedding -> FAISS top 30 -> cross-encoder rerank -> top 5 -> LLM
CANDIDATE_K = 30            # chunks pulled from FAISS before cross-encoder reranking
FINAL_K = 5                 # chunks handed to the LLM after reranking
CE_MIN_SCORE = 0.15         # sigmoid-normalized cross-encoder score floor; below this
                             # a chunk is treated as "not actually relevant" (tune against
                             # your own data - this is what now decides "no relevant docs found")

RECENCY_BOOST_WEIGHT = 0.05        # default small tie-break nudge toward newer docs
STRONG_RECENCY_BOOST_WEIGHT = 0.30  # applied for current/latest/management/designation queries

# Field names to look for when pulling a document's date out of chunk
# metadata. Add/reorder these to match your actual hbl_metadata.pkl schema.
DATE_FIELD_CANDIDATES = [
    "date", "doc_date", "document_date", "effective_date",
    "published_date", "publish_date", "last_updated",
    "updated_at", "version_date", "revision_date",
]

# Field names used to identify "the same underlying document" across
# versions (e.g. Tariff Schedule 2022 vs Tariff Schedule 2024). Adjust to
# match whatever unique-ish identifier your crawler/chunker attaches.
DOC_KEY_FIELD_CANDIDATES = [
    "doc_id", "document_id", "parent_doc_id",
    "source_file", "filename", "file_name", "source", "url", "title",
]

# Patterns stripped out when building a canonical doc key from a filename,
# so e.g. "kfs_savings_account_2023.pdf" and "kfs_savings_account_v2.pdf"
# collapse to the same key ("kfs savings account").
_VERSION_STRIP_PATTERNS = [
    r'[_\-\s]v(?:er(?:sion)?)?\.?\s*\d+(\.\d+)?',
    r'\((?:revised|updated|new|old|draft|superseded|archive[d]?)\)',
    r'[_\-\s](?:revised|updated|superseded|archived)',
    r'(19|20)\d{2}[-_/](0[1-9]|1[0-2])[-_/](0[1-9]|[12]\d|3[01])',
    r'(19|20)\d{2}[-_/](0[1-9]|1[0-2])',
    r'(19|20)\d{2}',
    r'\.(pdf|html|htm|docx?)$',
]

# Recurring HBL document series, matched against title/filename/url text.
# When a chunk's metadata matches one of these, it's grouped by the series
# name instead of the raw filename - this is what makes "Annual Report
# 2025" correctly outrank "Annual Report 2023/2021" even if the filenames
# don't share an obviously-strippable version suffix.
DOCUMENT_SERIES_PATTERNS = [
    (re.compile(r'annual report', re.I), 'series:annual_report'),
    (re.compile(r'\b(board of directors|board members)\b', re.I), 'series:board_of_directors'),
    (re.compile(r'\b(management (team|profile)|senior management|executive (team|management|committee))\b', re.I), 'series:management_profile'),
    (re.compile(r'\borganizational? (chart|structure)\b', re.I), 'series:org_structure'),
    (re.compile(r'\bcorporate profile\b', re.I), 'series:corporate_profile'),
    (re.compile(r'\bcode of conduct\b', re.I), 'series:code_of_conduct'),
    (re.compile(r'\bschedule of (charges|tariff)|tariff\b', re.I), 'series:tariff_schedule'),
    (re.compile(r'\bkey fact statement|\bkfs\b', re.I), 'series:kfs'),
]

YEAR_PATTERN = re.compile(r'(?<!\d)(19[5-9]\d|20[0-3]\d)(?!\d)')

HISTORICAL_KEYWORDS = [
    "old", "older", "oldest", "previous", "prior", "former",
    "earlier", "before", "used to", "historical", "history",
    "past", "outdated", "superseded", "archive", "archived",
    "original version", "back in", "as it was", "at the time",
    "then-current", "used to be",
]

RECENCY_SENSITIVE_KEYWORDS = [
    "current", "currently", "latest", "present", "presently",
    "today", "now", "as of", "existing", "sitting", "newest",
    "up to date", "up-to-date",
]

DESIGNATION_KEYWORDS = [
    "ceo", "chief executive officer", "cto", "chief technology officer",
    "cfo", "chief financial officer", "cro", "chief risk officer",
    "chro", "chief human resource officer", "chief human resources officer",
    "coo", "chief operating officer", "chairman", "chairperson", "president",
    "director", "board member", "board of directors", "head of",
    "designation", "management team", "senior management",
    "executive committee",
]

_PERSON_NAME_PATTERN = re.compile(
    r'\b(?:(?i:who\s+(?:is|was|leads|heads|runs|manages))|(?i:tell\s+me\s+about))\s+'
    r'([A-Z][\w.\'-]+(?:\s+[A-Z][\w.\'-]+){0,3})'
)


# ==================================================
# LOAD VECTOR DATABASE
# ==================================================

print("Loading FAISS index...")
index = faiss.read_index(INDEX_PATH)

with open(META_PATH, "rb") as f:
    chunks = pickle.load(f)

print("Chunks loaded:", len(chunks))


# ==================================================
# LOAD BGE-M3 (bi-encoder, for FAISS retrieval)
# ==================================================

device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading embedding model...")
embedder = SentenceTransformer(EMBED_MODEL, device=device)
embedder.max_seq_length = 1024

print("Embedding device:", device)


# ==================================================
# LOAD BGE-RERANKER-V2-M3 (cross-encoder, local only)
# ==================================================

print("Loading cross-encoder reranker...")
reranker = CrossEncoder(
    RERANKER_MODEL_PATH,
    max_length=512,
    device=device,
    trust_remote_code=True,
)

print("Reranker device:", device)


# ==================================================
# LOAD LOCAL LLM (Qwen2.5-0.5B-Instruct)
# ==================================================

llm_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

print("Loading LLM tokenizer...")
llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_PATH)

print("Loading LLM model...")
llm_model = AutoModelForCausalLM.from_pretrained(LLM_MODEL_PATH, torch_dtype=llm_dtype)
llm_model.to(device)
llm_model.eval()

print("LLM device:", device)


# ==================================================
# LOCAL LLM CHAT (replaces the Ollama HTTP call)
# ==================================================

def local_chat(messages, temperature=0, stream=False, max_new_tokens=MAX_NEW_TOKENS):
    prompt = llm_tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)

    do_sample = temperature > 0

    gen_kwargs = dict(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=do_sample,
        pad_token_id=llm_tokenizer.eos_token_id,
    )

    if do_sample:
        gen_kwargs["temperature"] = temperature
        gen_kwargs["top_p"] = 0.9
        gen_kwargs["repetition_penalty"] = 1.1

    if not stream:
        with torch.no_grad():
            output_ids = llm_model.generate(**gen_kwargs)
        new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
        return llm_tokenizer.decode(new_tokens, skip_special_tokens=True)

    streamer = TextIteratorStreamer(
        llm_tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )
    gen_kwargs["streamer"] = streamer

    thread = Thread(target=llm_model.generate, kwargs=gen_kwargs)
    thread.start()

    full_text = ""
    for token in streamer:
        print(token, end="", flush=True)
        full_text += token

    thread.join()
    print()
    return full_text


# ==================================================
# BANK ENTITY CHECK (word-boundary matching, not naive substring)
# ==================================================

HBL_TERMS = [
    "hbl",
    "habib bank limited",
    "habib bank ltd",
    "habib bank",
    "habib islamic",
]

OTHER_BANK_TERMS = [
    "meezan bank",
    "meezan",
    "ubl",
    "united bank limited",
    "united bank",
    "mcb bank",
    "mcb",
    "bank alfalah",
    "faysal bank",
    "askari bank",
    "allied bank limited",
    "allied bank",
    "abl",
    "bank al habib",
    "al habib",
    "js bank",
    "silkbank",
    "standard chartered",
    "soneri bank",
]


def _has_term(text, term):
    pattern = r"(?<!\w)" + re.escape(term) + r"(?!\w)"
    return re.search(pattern, text) is not None


def detect_bank_entity(query):
    q = query.lower()

    hbl_hit = any(_has_term(q, t) for t in HBL_TERMS)
    other_hit = any(_has_term(q, t) for t in OTHER_BANK_TERMS)

    if hbl_hit and other_hit:
        return "OTHER_BANK"

    if other_hit:
        return "OTHER_BANK"

    if hbl_hit:
        return "HBL"

    return "UNKNOWN"


# ==================================================
# PERSON / DESIGNATION / RECENCY QUERY DETECTION
# ==================================================

def is_person_or_designation_query(query):
    """Catches queries like 'Who is the CTO?' or 'Who is Jamal Nasir?' so
    they get routed straight to retrieval instead of risking a wrong
    block from the small classifier model."""
    q_lower = query.lower()
    if any(kw in q_lower for kw in DESIGNATION_KEYWORDS):
        return True
    if _PERSON_NAME_PATTERN.search(query):
        return True
    return False


def is_recency_sensitive_query(query):
    """Queries about current/latest info or management designations
    should strongly prefer the newest available document version."""
    q_lower = query.lower()
    if any(kw in q_lower for kw in RECENCY_SENSITIVE_KEYWORDS):
        return True
    if any(kw in q_lower for kw in DESIGNATION_KEYWORDS):
        return True
    return False


def extract_years_from_query(query):
    return [int(y) for y in YEAR_PATTERN.findall(query)]


def is_historical_query(query):
    q_lower = query.lower()
    if any(kw in q_lower for kw in HISTORICAL_KEYWORDS):
        return True
    if extract_years_from_query(query):
        return True
    return False


# ==================================================
# LLM SECURITY ROUTER
# ==================================================

KNOWN_CATEGORIES = [
    "HBL_RELATED",
    "OTHER_BANK",
    "WRITING_REQUEST",
    "CODING",
    "MATH",
    "MEDICAL",
    "AI_TECH",
    "GENERAL",
]

HARD_BLOCK_CATEGORIES = {
    "OTHER_BANK",
    "CODING",
    "MATH",
    "MEDICAL",
    "AI_TECH",
    "WRITING_REQUEST",
}


def classify_query(query, history=None):
    context_block = ""
    if history:
        lines = [f"{role}: {text}" for role, text in history]
        context_block = (
            "\nRecent conversation (context only, do not classify this):\n"
            + "\n".join(lines)
        )

    prompt = f"""
You are the security classifier for Habib Bank Limited AI Assistant.
Classify the CURRENT query into exactly ONE category.
Return only the category name, nothing else.

Default to HBL_RELATED or GENERAL whenever a query is ambiguous, short,
or could plausibly be about HBL (its staff, management, products,
services, or operations) - even if it does not say "HBL" explicitly.
Only choose a blocked category (OTHER_BANK, WRITING_REQUEST, CODING,
MATH, MEDICAL, AI_TECH) when the query is unambiguously and entirely
about that other topic, with no plausible banking angle.

Categories:
HBL_RELATED:
Anything directly or indirectly related to Habib Bank Limited - its
products, accounts, policies, staff, management, designations (CEO,
CTO, CFO, etc.), or general banking services - even without saying
"HBL" explicitly.
OTHER_BANK:
Query is clearly and specifically about a DIFFERENT, NAMED bank, with
no HBL connection.
WRITING_REQUEST:
Generic writing tasks unrelated to HBL (emails, CVs, letters, essays)
with no banking context at all.
CODING:
Programming/software questions with no banking context at all.
MATH:
Calculations, formulas, numerical problems unrelated to banking.
MEDICAL:
Health, diseases, medicines - with no banking context at all.
AI_TECH:
AI, models, LLMs, architecture - with no banking context at all.
GENERAL:
Ambiguous, unclear, or anything that doesn't obviously fit another
category. When in doubt, use this instead of a blocked category.
{context_block}
Current query:
{query}
Return only the category name.
"""

    raw = local_chat(
        [{"role": "user", "content": prompt}],
        temperature=0,
        stream=False,
        max_new_tokens=16,  # classification only needs one short word
    )
    result = raw.strip().upper()

    for cat in KNOWN_CATEGORIES:
        if cat in result:
            return cat

    return "GENERAL"


# ==================================================
# RECENCY-AWARE RERANKING HELPERS
# ==================================================

def parse_date_value(value):
    """Best-effort parse of a metadata value into a datetime.date."""
    if value is None:
        return None

    if isinstance(value, datetime.datetime):
        return value.date()
    if isinstance(value, datetime.date):
        return value

    if isinstance(value, (int, float)):
        try:
            return datetime.date(int(value), 1, 1)
        except (ValueError, OverflowError):
            return None

    if isinstance(value, str):
        text = value.strip()
        if not text:
            return None

        if _HAS_DATEUTIL:
            try:
                return _dateutil_parser.parse(text, fuzzy=True).date()
            except (ValueError, OverflowError, TypeError):
                pass

        m = re.search(r'(19|20)\d{2}[-/](0[1-9]|1[0-2])[-/](0[1-9]|[12]\d|3[01])', text)
        if m:
            try:
                return datetime.datetime.strptime(
                    m.group(0).replace('/', '-'), "%Y-%m-%d"
                ).date()
            except ValueError:
                pass

        m = re.search(r'(19|20)\d{2}', text)
        if m:
            return datetime.date(int(m.group(0)), 1, 1)

    return None


def get_doc_date(meta):
    """Pull a usable date out of chunk metadata, trying explicit date
    fields first, then falling back to parsing filename/url/title."""
    for field in DATE_FIELD_CANDIDATES:
        if meta.get(field):
            parsed = parse_date_value(meta[field])
            if parsed:
                return parsed

    for field in ("source_file", "filename", "file_name", "source", "url", "title"):
        if meta.get(field):
            parsed = parse_date_value(str(meta[field]))
            if parsed:
                return parsed

    return None


def detect_doc_series(meta):
    """Match against recurring HBL document series (annual report, board
    of directors, tariff schedule, etc.) regardless of filename quirks."""
    text_fields = []
    for field in ("title", "source_file", "filename", "file_name", "source", "url"):
        if meta.get(field):
            text_fields.append(str(meta[field]))
    combined = " ".join(text_fields)

    for pattern, series in DOCUMENT_SERIES_PATTERNS:
        if pattern.search(combined):
            return series
    return None


def canonicalize_doc_key(meta):
    """Build a version-agnostic key so different dated/versioned copies
    of the same underlying document (annual report, tariff sheet, KFS,
    circular, etc.) group together. Prefers a known document series match
    (more robust) and falls back to stripping version/date patterns off
    whatever filename-ish field is available."""
    series = detect_doc_series(meta)
    if series:
        return series

    raw = None
    for field in DOC_KEY_FIELD_CANDIDATES:
        if meta.get(field):
            raw = str(meta[field])
            break

    if not raw:
        return None

    key = raw.lower().strip()
    for pattern in _VERSION_STRIP_PATTERNS:
        key = re.sub(pattern, ' ', key, flags=re.IGNORECASE)
    key = re.sub(r'[\W_]+', ' ', key).strip()

    return key or None


def _rank_latest_only(candidates, weight=RECENCY_BOOST_WEIGHT):
    """Collapse multiple versions of the same document down to the
    latest dated version, then rank survivors by (cross-encoder)
    relevance with a recency tie-breaker. `weight` controls how strongly
    recency is allowed to matter - pass STRONG_RECENCY_BOOST_WEIGHT for
    current/latest/management/designation queries."""

    latest_date_by_key = {}
    for c in candidates:
        key, date = c["doc_key"], c["date"]
        if key is None or date is None:
            continue
        if key not in latest_date_by_key or date > latest_date_by_key[key]:
            latest_date_by_key[key] = date

    survivors = []
    for c in candidates:
        key, date = c["doc_key"], c["date"]
        if key is not None and key in latest_date_by_key:
            if date is not None and date < latest_date_by_key[key]:
                continue  # older version of a doc we have a newer copy of
            if date is None:
                continue  # undated chunk from a doc that has dated siblings
        survivors.append(c)

    known_dates = [c["date"] for c in survivors if c["date"] is not None]
    newest = max(known_dates) if known_dates else None
    oldest = min(known_dates) if known_dates else None
    span_days = (newest - oldest).days if (newest and oldest and newest != oldest) else 0

    def final_score(c):
        if c["date"] is not None and span_days > 0:
            recency = (c["date"] - oldest).days / span_days
        elif c["date"] is not None:
            recency = 1.0
        else:
            recency = -0.02  # tiny penalty for unknown date, not exclusion
        return c["score"] + weight * recency

    survivors.sort(key=final_score, reverse=True)
    return [c["meta"] for c in survivors[:FINAL_K]]


def _rank_historical(candidates, requested_years):
    """Keep every version in play. If the user named a specific year,
    surface matching chunks first; otherwise rank by relevance across
    all versions."""

    def year_match(c):
        if not requested_years or c["date"] is None:
            return False
        return c["date"].year in requested_years

    candidates.sort(key=lambda c: (year_match(c), c["score"]), reverse=True)
    return [c["meta"] for c in candidates[:FINAL_K]]


# ==================================================
# RETRIEVAL
# Query -> BGE-M3 embedding -> FAISS top 30 -> cross-encoder rerank -> top 5
# ==================================================

def retrieve(query):
    vector = embedder.encode([query], normalize_embeddings=True)
    vector = vector.astype("float32")

    scores, ids = index.search(vector, CANDIDATE_K)

    pool = []
    for score, idx in zip(scores[0], ids[0]):
        if idx == -1:
            continue
        if score < MIN_SCORE:
            continue
        pool.append((float(score), chunks[idx]))

    if not pool:
        return []

    # --- cross-encoder rerank -------------------------------------
    pairs = [(query, meta["text"]) for _, meta in pool]
    raw_ce_scores = reranker.predict(pairs)
    ce_scores = 1 / (1 + np.exp(-np.asarray(raw_ce_scores, dtype=float)))  # sigmoid -> 0..1

    candidates = []
    for (bi_score, meta), ce_score in zip(pool, ce_scores):
        if ce_score < CE_MIN_SCORE:
            continue  # cross-encoder says this really isn't relevant
        candidates.append({
            "meta": meta,
            "score": float(ce_score),
            "bi_score": bi_score,
            "doc_key": canonicalize_doc_key(meta),
            "date": get_doc_date(meta),
        })

    if not candidates:
        return []

    # --- recency-aware final selection ------------------------------
    if is_historical_query(query):
        return _rank_historical(candidates, extract_years_from_query(query))

    weight = STRONG_RECENCY_BOOST_WEIGHT if is_recency_sensitive_query(query) else RECENCY_BOOST_WEIGHT
    return _rank_latest_only(candidates, weight=weight)


# ==================================================
# GENERATION
# ==================================================

SYSTEM_PROMPT = (
    "You are Habib Bank Limited Assistant. Answer only from retrieved HBL "
    "documents. Do not guess. If unavailable, say you could not find this "
    "information."
)


def generate_answer(query, docs):
    context = "\n\n".join(d["text"] for d in docs)

    return local_chat(
        [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": f"""
HBL DOCUMENTS:
{context}
QUESTION:
{query}
""",
            },
        ],
        temperature=0.1,
        stream=True,
    )


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

print("\n" + "=" * 60)
print("HBL Assistant Ready")
print("Type exit to quit")
print("=" * 60)

history = []  # list of (role, text) tuples, most recent last

while True:
    query = input("\nUser: ")

    if query.lower() == "exit":
        break

    if not query.strip():
        continue

    # ----------------------------------------------
    # ENTITY CHECK FIRST (cheap, catches obvious cases)
    # ----------------------------------------------
    entity = detect_bank_entity(query)

    if entity == "OTHER_BANK":
        print("\nAssistant:\nI can only assist with Habib Bank Limited related queries.")
        continue

    if entity == "HBL":
        intent = "HBL_RELATED"
    elif is_person_or_designation_query(query):
        # "Who is the CTO?" / "Who is Jamal Nasir?" - no bank named, but
        # clearly a person/designation lookup. Let retrieval decide
        # relevance instead of risking a wrong block from the classifier.
        intent = "HBL_RELATED"
    else:
        intent = classify_query(query, history)

    print("Intent:", intent)

    if intent in HARD_BLOCK_CATEGORIES:
        print("\nAssistant:\nI can only assist with Habib Bank Limited related queries.")
        continue

    # ----------------------------------------------
    # RETRIEVE (this is the real relevance gate)
    # ----------------------------------------------
    docs = retrieve(query)

    if len(docs) == 0:
        print("\nAssistant:\nI could not find this information in the available HBL documents.")
        continue

    # ----------------------------------------------
    # GENERATE
    # ----------------------------------------------
    print("\nAssistant:")
    answer = generate_answer(query, docs)

    history.append(("user", query))
    history.append(("assistant", answer))
    history = history[-HISTORY_TURNS * 2:]