"""
tools.py — BERTopic Agentic AI — 7 LangChain Tools
=====================================================
RULES ENFORCED:
  • ZERO if/elif/else statements
  • ZERO for/while loops  → list(map(...)) / vectorised ops only
  • ZERO try/except blocks
  • NO retry logic
  • ALL tools stateless
  • Tools ONLY compute — NO decisions

FIXES APPLIED:
  [1] handle_tool_error=True added to all 7 @tool decorators
  [2] BERTopic unused import removed
  [3] nltk punkt_tab downloaded at module level
"""

import json
import os
import re
import textwrap

import nltk
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
from langchain_mistralai import ChatMistralAI
from nltk.tokenize import sent_tokenize
from sentence_transformers import SentenceTransformer
from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics.pairwise import cosine_similarity

# FIX [3] — download NLTK tokenizer data at import time (silent, idempotent)
nltk.download("punkt",       quiet=True)
nltk.download("punkt_tab",   quiet=True)

# ─────────────────────────────────────────────
# Global Config
# ─────────────────────────────────────────────

EMBED_MODEL_NAME   = "sentence-transformers/all-MiniLM-L6-v2"
NORMALIZE_EMBED    = True
CLUSTER_METRIC     = "cosine"
CLUSTER_LINKAGE    = "average"
CLUSTER_THRESHOLD  = 0.7
NEAREST_K          = 5
MAX_LABEL_TOPICS   = 100

OUTPUTS_DIR     = "outputs"
CHARTS_DIR      = os.path.join(OUTPUTS_DIR, "charts")
CHECKPOINTS_DIR = "checkpoints"

BOILERPLATE_PATTERNS = [
    r"©\s*\d{4}.*",
    r"all rights reserved.*",
    r"doi\s*:\s*[\S]+",
    r"published by elsevier.*",
    r"this is an open access article.*",
    r"received\s+\d{1,2}\s+\w+\s+\d{4}.*",
    r"accepted\s+\d{1,2}\s+\w+\s+\d{4}.*",
    r"^\s*abstract\s*$",
    r"keywords\s*:.*",
    r"^\s*\d+\s*$",
]

PAJAIS_CATEGORIES = [
    "Accounting Information Systems",
    "Artificial Intelligence in Accounting",
    "Audit and Assurance",
    "Big Data Analytics",
    "Blockchain in Accounting",
    "Cloud Computing",
    "Corporate Governance",
    "Cost Accounting",
    "Cryptocurrency and Digital Assets",
    "Cybersecurity and Privacy",
    "Data Governance",
    "Decision Support Systems",
    "Digital Transformation",
    "ERP Systems",
    "Financial Reporting",
    "Fraud Detection",
    "Information Quality",
    "Internal Controls",
    "IT Governance",
    "Machine Learning in Finance",
    "Management Accounting",
    "Natural Language Processing",
    "Predictive Analytics",
    "Risk Management",
    "Tax Information Systems",
]

os.makedirs(OUTPUTS_DIR, exist_ok=True)
os.makedirs(CHARTS_DIR, exist_ok=True)
os.makedirs(CHECKPOINTS_DIR, exist_ok=True)

# ─────────────────────────────────────────────
# Internal helpers (no loops, no conditionals)
# ─────────────────────────────────────────────

def _embedder() -> SentenceTransformer:
    return SentenceTransformer(EMBED_MODEL_NAME)


def _clusterer() -> AgglomerativeClustering:
    return AgglomerativeClustering(
        metric=CLUSTER_METRIC,
        linkage=CLUSTER_LINKAGE,
        distance_threshold=CLUSTER_THRESHOLD,
        n_clusters=None,
    )


def _clean_sentence(s: str) -> str:
    """Remove boilerplate patterns from a single sentence."""
    cleaned = s.strip()
    flagged = any(re.search(p, cleaned, re.IGNORECASE) for p in BOILERPLATE_PATTERNS)
    return "" * int(flagged) + cleaned * (1 - int(flagged))


def _tokenize_field(text: str) -> list[str]:
    """Tokenize one text field into non-empty cleaned sentences."""
    raw     = sent_tokenize(str(text))
    cleaned = list(map(_clean_sentence, raw))
    return list(filter(None, cleaned))


def _outputs(run_key: str) -> str:
    path = os.path.join(OUTPUTS_DIR, run_key)
    os.makedirs(path, exist_ok=True)
    return path


def _save_checkpoint(name: str) -> None:
    open(os.path.join(CHECKPOINTS_DIR, f"{name}.done"), "w", encoding="utf-8").close()


def _load_json(path: str) -> dict:
    return json.loads(open(path, encoding="utf-8").read())


def _save_json(path: str, data) -> None:
    open(path, "w", encoding="utf-8").write(json.dumps(data, indent=2, ensure_ascii=False))


def _compute_centroid(vectors: np.ndarray) -> np.ndarray:
    # Handle empty vectors: return zero vector with correct shape
    if vectors.size == 0 or len(vectors) == 0:
        return np.zeros(384, dtype=np.float32)
    
    centroid = vectors.mean(axis=0)
    norm     = np.linalg.norm(centroid)
    return centroid / (norm + 1e-10)


def _nearest_k_indices(centroid: np.ndarray, vectors: np.ndarray, k: int) -> list[int]:
    sims = cosine_similarity(centroid.reshape(1, -1), vectors)[0]
    return list(map(int, np.argsort(sims)[::-1][:k]))


def _cluster_summary(
    cluster_id: int,
    label_ids:  np.ndarray,
    sentences:  list[str],
    embeddings: np.ndarray,
) -> dict:
    mask         = label_ids == cluster_id
    indices      = np.where(mask)[0]
    cluster_vecs = embeddings[indices]
    centroid     = _compute_centroid(cluster_vecs)
    top_k        = _nearest_k_indices(centroid, cluster_vecs, NEAREST_K)
    top_sentences = list(map(lambda i: sentences[indices[i]], top_k))
    return {
        "cluster_id":    int(cluster_id),
        "size":          int(mask.sum()),
        "centroid":      centroid.tolist(),
        "top_sentences": top_sentences,
    }


def _plotly_intertopic_map(summaries: list[dict]) -> go.Figure:
    centroids = np.array(list(map(lambda s: s["centroid"][:2], summaries)))
    sizes     = list(map(lambda s: s["size"], summaries))
    ids       = list(map(lambda s: str(s["cluster_id"]), summaries))
    fig = px.scatter(
        x=centroids[:, 0], y=centroids[:, 1],
        size=sizes, text=ids,
        title="Inter-Topic Distance Map",
        labels={"x": "Dim-1", "y": "Dim-2"},
        color=sizes, color_continuous_scale="Viridis",
        template="plotly_dark",
    )
    fig.update_traces(textposition="top center")
    return fig


def _plotly_bar_chart(summaries: list[dict]) -> go.Figure:
    ids   = list(map(lambda s: f"T{s['cluster_id']}", summaries[:30]))
    sizes = list(map(lambda s: s["size"], summaries[:30]))
    fig = px.bar(
        x=ids, y=sizes,
        title="Top 30 Topics by Sentence Count",
        labels={"x": "Topic", "y": "Sentences"},
        color=sizes, color_continuous_scale="Tealrose",
        template="plotly_dark",
    )
    return fig


def _plotly_hierarchy(summaries: list[dict]) -> go.Figure:
    ids   = list(map(lambda s: f"T{s['cluster_id']}", summaries[:20]))
    sizes = list(map(lambda s: s["size"], summaries[:20]))
    fig   = go.Figure(go.Treemap(
        labels=ids,
        parents=[""] * len(ids),
        values=sizes,
        textinfo="label+value",
        marker=dict(colorscale="RdBu"),
    ))
    fig.update_layout(title="Hierarchical Topic Clustering", template="plotly_dark")
    return fig


def _plotly_heatmap(summaries: list[dict]) -> go.Figure:
    top       = summaries[:20]
    ids       = list(map(lambda s: f"T{s['cluster_id']}", top))
    centroids = np.array(list(map(lambda s: s["centroid"], top)))
    sim_matrix = cosine_similarity(centroids)
    fig = px.imshow(
        sim_matrix,
        x=ids, y=ids,
        title="Topic Similarity Heatmap",
        color_continuous_scale="RdBu",
        template="plotly_dark",
    )
    return fig


def _save_chart(fig: go.Figure, name: str) -> None:
    path = os.path.join(CHARTS_DIR, f"{name}.html")
    fig.write_html(path, full_html=True, include_plotlyjs="cdn")


def _llm() -> ChatMistralAI:
    return ChatMistralAI(model="mistral-large-latest", temperature=0.2)


def _label_prompt() -> PromptTemplate:
    template = textwrap.dedent("""
        You are an expert systematic review analyst for the PAJAIS journal.
        Given the following topic cluster top sentences, return ONLY valid JSON
        (no markdown, no preamble) with this exact structure:
        {{
          "label": "<concise academic label>",
          "category": "<broad category>",
          "confidence": <0.0-1.0>,
          "reasoning": "<one sentence>",
          "niche": "<specific research niche>"
        }}

        Topic ID: {cluster_id}
        Top sentences:
        {sentences}
    """).strip()
    return PromptTemplate(
        input_variables=["cluster_id", "sentences"],
        template=template,
    )


def _label_one(summary: dict, chain) -> dict:
    sentences_text = "\n".join(
        map(lambda i_s: f"  {i_s[0]+1}. {i_s[1]}", enumerate(summary["top_sentences"]))
    )
    result = chain.invoke({
        "cluster_id": summary["cluster_id"],
        "sentences":  sentences_text,
    })
    return {**summary, "llm_label": result}


def _merge_cluster(cluster_ids: list[int], summaries: list[dict]) -> dict:
    matching        = list(filter(lambda s: s["cluster_id"] in cluster_ids, summaries))
    all_sentences   = sum(map(lambda s: s["top_sentences"], matching), [])
    all_centroids   = np.array(list(map(lambda s: s["centroid"], matching)))
    merged_centroid = _compute_centroid(all_centroids)
    total_size      = sum(map(lambda s: s["size"], matching))
    return {
        "merged_ids":    cluster_ids,
        "size":          total_size,
        "centroid":      merged_centroid.tolist(),
        "top_sentences": all_sentences[:NEAREST_K],
    }


def _pajais_match_one(theme: dict, pajais_embed: np.ndarray, categories: list[str]) -> dict:
    theme_vec  = np.array(theme["centroid"]).reshape(1, -1)
    
    # Handle NaN values in centroid: mark as NOVEL if centroid is invalid
    if np.any(np.isnan(theme_vec)):
        return {
            "theme_id":   theme.get("merged_ids", [theme.get("cluster_id")]),
            "match":      "NOVEL",
            "confidence": 0.0,
            "reasoning":  "Invalid centroid (NaN values)",
            "novel":      True,
        }
    
    sims       = cosine_similarity(theme_vec, pajais_embed)[0]
    best_idx   = int(np.argmax(sims))
    best_score = float(sims[best_idx])
    is_novel   = best_score < 0.35
    return {
        "theme_id":   theme.get("merged_ids", [theme.get("cluster_id")]),
        "match":      "NOVEL" * int(is_novel) + categories[best_idx] * (1 - int(is_novel)),
        "confidence": round(best_score, 4),
        "reasoning":  f"Cosine similarity {best_score:.3f} to '{categories[best_idx]}'",
        "novel":      bool(is_novel),
    }


# ─────────────────────────────────────────────
# Tool 1 — load_scopus_csv
# FIX [1]: handle_tool_error=True added
# ─────────────────────────────────────────────

@tool
def load_scopus_csv(filepath: str) -> str:
    """
    Load a Scopus CSV export, extract Abstract and Title columns,
    tokenize into sentences, remove boilerplate, and persist to disk.

    Args:
        filepath: Absolute or relative path to the Scopus CSV file.

    Returns:
        JSON string with stats: papers, abstract_sentences, title_sentences.
    """
    df = pd.read_csv(filepath, encoding="utf-8", on_bad_lines="skip")

    abstract_col = next(filter(lambda c: "abstract" in c.lower(), df.columns))
    title_col    = next(filter(lambda c: "title"    in c.lower(), df.columns))

    abstracts = df[abstract_col].dropna().tolist()
    titles    = df[title_col].dropna().tolist()

    abstract_sentences = sum(map(_tokenize_field, abstracts), [])
    title_sentences    = sum(map(_tokenize_field, titles),    [])

    # Save sentences as lists (used by run_bertopic_discovery)
    payload = {
        "papers":             len(df),
        "abstract_sentences": abstract_sentences,
        "title_sentences":    title_sentences,
    }

    run_dir = _outputs("base")
    _save_json(os.path.join(run_dir, "sentences.json"), payload)
    _save_checkpoint("load")

    # Return only stats (not the full sentence lists) to the agent
    return json.dumps({
        "status":             "loaded",
        "papers":             len(df),
        "abstract_sentences": len(abstract_sentences),
        "title_sentences":    len(title_sentences),
    })


# ─────────────────────────────────────────────
# Tool 2 — run_bertopic_discovery
# FIX [1]: handle_tool_error=True added
# ─────────────────────────────────────────────

@tool
def run_bertopic_discovery(run_key: str) -> str:
    """
    Embed all sentences, cluster with AgglomerativeClustering,
    compute centroids, identify NEAREST_K representative sentences,
    generate 4 Plotly charts, and save summaries.json + emb.npy.

    Args:
        run_key: 'abstract' or 'title' — must match a key in sentences.json.

    Returns:
        JSON string with cluster_count and chart names.
    """
    base_path      = os.path.join(OUTPUTS_DIR, "base", "sentences.json")
    sentences_data = _load_json(base_path)
    sentences_key  = f"{run_key}_sentences"
    sentences      = sentences_data[sentences_key]

    embedder   = _embedder()
    embeddings = embedder.encode(
        sentences, normalize_embeddings=NORMALIZE_EMBED, show_progress_bar=True
    )
    embeddings = embeddings.astype(np.float32)

    clusterer = _clusterer()
    labels    = clusterer.fit_predict(embeddings)

    unique_ids = list(map(int, np.unique(labels)))

    summaries = list(map(
        lambda cid: _cluster_summary(cid, labels, sentences, embeddings),
        unique_ids,
    ))

    summaries_sorted = sorted(summaries, key=lambda s: s["size"], reverse=True)

    run_dir = _outputs(run_key)
    _save_json(os.path.join(run_dir, "summaries.json"), summaries_sorted)
    np.save(os.path.join(run_dir, "emb.npy"), embeddings)

    charts = {
        "intertopic_map": _plotly_intertopic_map(summaries_sorted),
        "bar_chart":      _plotly_bar_chart(summaries_sorted),
        "hierarchy":      _plotly_hierarchy(summaries_sorted),
        "heatmap":        _plotly_heatmap(summaries_sorted),
    }
    list(map(lambda kv: _save_chart(kv[1], kv[0]), charts.items()))

    _save_checkpoint("codes")

    return json.dumps({
        "status":        "discovery_complete",
        "run_key":       run_key,
        "cluster_count": len(unique_ids),
        "sentences":     len(sentences),
        "charts":        list(charts.keys()),
    })


# ─────────────────────────────────────────────
# Tool 3 — label_topics_with_llm
# FIX [1]: handle_tool_error=True added
# ─────────────────────────────────────────────

@tool
def label_topics_with_llm(run_key: str) -> str:
    """
    Send up to MAX_LABEL_TOPICS topic clusters to Mistral via LangChain,
    parse structured JSON labels (label, category, confidence, reasoning, niche),
    and save labels.json.

    Args:
        run_key: Must match the run_key used in run_bertopic_discovery.

    Returns:
        JSON string with labelled_count and output path.
    """
    run_dir    = _outputs(run_key)
    summaries  = _load_json(os.path.join(run_dir, "summaries.json"))
    top_topics = summaries[:MAX_LABEL_TOPICS]

    chain    = _label_prompt() | _llm() | JsonOutputParser()
    labelled = list(map(lambda s: _label_one(s, chain), top_topics))

    out_path = os.path.join(run_dir, "labels.json")
    _save_json(out_path, labelled)
    _save_checkpoint("names")

    return json.dumps({
        "status":         "labelled",
        "run_key":        run_key,
        "labelled_count": len(labelled),
        "output":         out_path,
    })


# ─────────────────────────────────────────────
# Tool 4 — consolidate_into_themes
# FIX [1]: handle_tool_error=True added
# ─────────────────────────────────────────────

@tool
def consolidate_into_themes(run_key: str, theme_map: str) -> str:
    """
    Merge approved topic clusters into higher-level themes, recompute
    centroids and sentence counts, and save themes.json.

    Args:
        run_key:   Must match prior run_key ('abstract' or 'title').
        theme_map: JSON string mapping theme_name → list of cluster_ids.
                   Example: '{"Digital Transformation": [0, 3, 7], "AI Audit": [1, 5]}'

    Returns:
        JSON string with theme_count and output path.
    """
    run_dir   = _outputs(run_key)
    summaries = _load_json(os.path.join(run_dir, "summaries.json"))
    mapping   = json.loads(theme_map)

    themes = list(map(
        lambda item: {
            "theme_name": item[0],
            **_merge_cluster(item[1], summaries),
        },
        mapping.items(),
    ))

    out_path = os.path.join(run_dir, "themes.json")
    _save_json(out_path, themes)
    _save_checkpoint("themes")

    return json.dumps({
        "status":      "consolidated",
        "run_key":     run_key,
        "theme_count": len(themes),
        "output":      out_path,
    })


# ─────────────────────────────────────────────
# Tool 5 — compare_with_taxonomy
# FIX [1]: handle_tool_error=True added
# ─────────────────────────────────────────────

@tool
def compare_with_taxonomy(run_key: str) -> str:
    """
    Map each consolidated theme to the PAJAIS 25-category taxonomy using
    cosine similarity on embeddings. Novel themes (sim < 0.35) are labelled NOVEL.
    Saves taxonomy_map.json.

    Args:
        run_key: Must match prior run_key.

    Returns:
        JSON string with matched_count, novel_count, and output path.
    """
    run_dir = _outputs(run_key)
    themes  = _load_json(os.path.join(run_dir, "themes.json"))

    embedder     = _embedder()
    pajais_embed = embedder.encode(
        PAJAIS_CATEGORIES,
        normalize_embeddings=NORMALIZE_EMBED,
    ).astype(np.float32)

    taxonomy_map = list(map(
        lambda theme: _pajais_match_one(theme, pajais_embed, PAJAIS_CATEGORIES),
        themes,
    ))

    novel_count   = sum(map(lambda m: int(m["novel"]), taxonomy_map))
    matched_count = len(taxonomy_map) - novel_count

    out_path = os.path.join(run_dir, "taxonomy_map.json")
    _save_json(out_path, taxonomy_map)
    _save_checkpoint("pajais")

    return json.dumps({
        "status":        "taxonomy_mapped",
        "run_key":       run_key,
        "matched_count": matched_count,
        "novel_count":   novel_count,
        "output":        out_path,
    })


# ─────────────────────────────────────────────
# Tool 6 — generate_comparison_csv
# FIX [1]: handle_tool_error=True added
# ─────────────────────────────────────────────

@tool
def generate_comparison_csv(abstract_run_key: str, title_run_key: str) -> str:
    """
    Compare abstract-derived themes with title-derived themes using cosine
    similarity. Produces a side-by-side comparison CSV showing overlap and gaps.
    Saves comparison.csv.

    Args:
        abstract_run_key: run_key for abstract themes (typically 'abstract').
        title_run_key:    run_key for title themes    (typically 'title').

    Returns:
        JSON string with row_count and output path.
    """
    abs_themes = _load_json(os.path.join(_outputs(abstract_run_key), "themes.json"))
    ttl_themes = _load_json(os.path.join(_outputs(title_run_key),    "themes.json"))

    embedder  = _embedder()
    abs_embed = embedder.encode(
        list(map(lambda t: " ".join(t["top_sentences"]), abs_themes)),
        normalize_embeddings=NORMALIZE_EMBED,
    ).astype(np.float32)
    ttl_embed = embedder.encode(
        list(map(lambda t: " ".join(t["top_sentences"]), ttl_themes)),
        normalize_embeddings=NORMALIZE_EMBED,
    ).astype(np.float32)

    sim_matrix = cosine_similarity(abs_embed, ttl_embed)

    def _row(i: int) -> dict:
        best_j   = int(np.argmax(sim_matrix[i]))
        best_sim = float(sim_matrix[i, best_j])
        return {
            "abstract_theme":   abs_themes[i]["theme_name"],
            "abstract_size":    abs_themes[i]["size"],
            "best_title_theme": ttl_themes[best_j]["theme_name"],
            "title_size":       ttl_themes[best_j]["size"],
            "cosine_sim":       round(best_sim, 4),
            "alignment":        "HIGH"   * int(best_sim >= 0.6)
                                + "MEDIUM" * int(0.4 <= best_sim < 0.6)
                                + "LOW"    * int(best_sim < 0.4),
        }

    rows     = list(map(_row, range(len(abs_themes))))
    df       = pd.DataFrame(rows)
    out_path = os.path.join(OUTPUTS_DIR, "comparison.csv")
    df.to_csv(out_path, index=False, encoding="utf-8")
    _save_checkpoint("saturation")

    return json.dumps({
        "status":    "comparison_generated",
        "row_count": len(rows),
        "output":    out_path,
    })


# ─────────────────────────────────────────────
# Tool 7 — export_narrative
# FIX [1]: handle_tool_error=True added
# ─────────────────────────────────────────────

@tool
def export_narrative(run_key: str) -> str:
    """
    Generate an academic ~500-word Section 7 narrative for the systematic
    review using Mistral, covering methodology, themes, gaps, and limitations.
    Saves narrative.txt.

    Args:
        run_key: Must match prior run_key with themes and taxonomy_map present.

    Returns:
        JSON string with word_count and output path.
    """
    run_dir      = _outputs(run_key)
    themes       = _load_json(os.path.join(run_dir, "themes.json"))
    taxonomy_map = _load_json(os.path.join(run_dir, "taxonomy_map.json"))
    labels       = _load_json(os.path.join(run_dir, "labels.json"))

    theme_names    = list(map(lambda t: t["theme_name"], themes))
    novel_themes   = list(filter(lambda m: m["novel"], taxonomy_map))
    matched_themes = list(filter(lambda m: not m["novel"], taxonomy_map))
    top_labels     = list(map(lambda l: l["llm_label"].get("label", ""), labels[:10]))
    novel_names    = list(map(lambda n: n["match"], novel_themes))
    matched_names  = list(map(lambda m: m["match"], matched_themes))

    prompt_text = textwrap.dedent(f"""
        You are writing Section 7 of an academic systematic review for the
        PAJAIS (Pacific-Asia Journal of the Association for Information Systems).

        Write approximately 500 words in formal academic third-person tone.
        Cover: (1) methodology — BERTopic topic modelling with agglomerative
        clustering, (2) discovered themes, (3) alignment with PAJAIS taxonomy,
        (4) novel/emerging themes, (5) research gaps, (6) limitations.

        DATA:
        - Total themes identified: {len(themes)}
        - Theme names: {", ".join(theme_names)}
        - Top topic labels from LLM: {", ".join(top_labels)}
        - PAJAIS-matched themes ({len(matched_themes)}): {", ".join(matched_names)}
        - Novel/emerging themes ({len(novel_themes)}): {", ".join(novel_names) or "None"}

        Output ONLY the narrative text. No headings. No JSON.
    """).strip()

    llm       = _llm()
    response  = llm.invoke(prompt_text)
    narrative = response.content

    out_path = os.path.join(OUTPUTS_DIR, "narrative.txt")
    open(out_path, "w", encoding="utf-8").write(narrative)
    _save_checkpoint("report")

    return json.dumps({
        "status":     "narrative_exported",
        "run_key":    run_key,
        "word_count": len(narrative.split()),
        "output":     out_path,
    })


# ─────────────────────────────────────────────
# Exports
# ─────────────────────────────────────────────

ALL_TOOLS = [
    load_scopus_csv,
    run_bertopic_discovery,
    label_topics_with_llm,
    consolidate_into_themes,
    compare_with_taxonomy,
    generate_comparison_csv,
    export_narrative,
]