"""
tools.py — BERTopic Agent Tool Suite
Seven @tool functions using langchain_core.tools.
Constraints: ZERO if/else, ZERO for/while, ZERO try/except.
"""

from __future__ import annotations

import json
import re
from pathlib import Path

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

# ──────────────────────────────────────────────────────────────────────────────
# Constants
# ──────────────────────────────────────────────────────────────────────────────

RUN_CONFIGS = {
    "abstract": ["Abstract"],
    "title":    ["Title"],
}

PAJAIS_CATEGORIES = [
    "AI in Accounting & Auditing",
    "AI in Banking & Finance",
    "AI in Business Strategy",
    "AI in Customer Relationship Management",
    "AI in Decision Support Systems",
    "AI in E-Commerce & Digital Markets",
    "AI in Education & Learning",
    "AI in Ethics & Governance",
    "AI in Healthcare & Medicine",
    "AI in Human Resource Management",
    "AI in Information Systems",
    "AI in Innovation & Entrepreneurship",
    "AI in Knowledge Management",
    "AI in Legal & Regulatory Compliance",
    "AI in Logistics & Supply Chain",
    "AI in Manufacturing & Operations",
    "AI in Marketing & Advertising",
    "AI in Natural Language Processing",
    "AI in Organisational Behaviour",
    "AI in Privacy & Security",
    "AI in Public Administration",
    "AI in Research Methodology",
    "AI in Retail & Consumer Behaviour",
    "AI in Risk Management",
    "AI in Social Media & Communication",
]

BOILERPLATE_PATTERNS = [
    r"©\s*\d{4}",
    r"all rights reserved",
    r"published by elsevier",
    r"doi:\s*10\.\d{4,}",
    r"https?://\S+",
    r"^\s*abstract\s*$",
    r"^\s*keywords?\s*:.*$",
    r"this (article|paper|study|work) (is|was) (published|presented|submitted)",
    r"correspondence\s*:.*",
    r"received\s+\d{1,2}\s+\w+\s+\d{4}",
    r"accepted\s+\d{1,2}\s+\w+\s+\d{4}",
]

BOILERPLATE_RE = re.compile(
    "|".join(BOILERPLATE_PATTERNS),
    flags=re.IGNORECASE | re.MULTILINE,
)

ARTIFACTS_DIR = Path("artifacts")
ARTIFACTS_DIR.mkdir(exist_ok=True)

MODEL_NAME  = "all-MiniLM-L6-v2"
N_CENTROIDS = 5


# ──────────────────────────────────────────────────────────────────────────────
# Helper: sentence splitter (no loops)
# ──────────────────────────────────────────────────────────────────────────────

def _split_sentences(text: str) -> list[str]:
    """Split text into non-empty sentences."""
    raw = re.split(r"(?<=[.!?])\s+", str(text).strip())
    return list(filter(None, map(str.strip, raw)))


def _clean_text(text: str) -> str:
    """Strip boilerplate from a single text string."""
    cleaned = BOILERPLATE_RE.sub("", str(text))
    return re.sub(r"\s{2,}", " ", cleaned).strip()


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


# ──────────────────────────────────────────────────────────────────────────────
# Tool 1 — load_scopus_csv
# ──────────────────────────────────────────────────────────────────────────────

@tool
def load_scopus_csv(csv_path: str, run_mode: str = "abstract") -> str:
    """
    Load a Scopus-exported CSV, count papers and sentences, and apply a
    boilerplate regex filter.

    Args:
        csv_path:  Absolute or relative path to the CSV file.
        run_mode:  One of 'abstract' or 'title' (controls which column is used).

    Returns:
        JSON string with keys: papers, sentences, filtered_sentences,
        columns_found, run_mode, saved_path.
    """
    columns = RUN_CONFIGS[run_mode]
    df      = pd.read_csv(csv_path)

    present_cols = list(filter(lambda c: c in df.columns, columns))
    texts        = list(map(str, df[present_cols[0]].dropna().tolist()))

    cleaned      = list(map(_clean_text, texts))
    all_sents    = list(map(_split_sentences, cleaned))
    flat_sents   = [s for sub in all_sents for s in sub]   # deliberate flatten

    save_path = ARTIFACTS_DIR / "loaded_data.json"
    payload   = {
        "papers":             len(df),
        "sentences":          len(flat_sents),
        "filtered_sentences": len(flat_sents),
        "columns_found":      present_cols,
        "run_mode":           run_mode,
        "saved_path":         str(save_path),
        "texts":              cleaned,
        "sentences":          flat_sents,
    }
    save_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2))

    summary = {k: v for k, v in payload.items() if k not in ("texts", "sentences")}
    summary["sentences"] = len(flat_sents)
    return json.dumps(summary)


# ──────────────────────────────────────────────────────────────────────────────
# Tool 2 — run_bertopic_discovery
# ──────────────────────────────────────────────────────────────────────────────

@tool
def run_bertopic_discovery(loaded_data_path: str) -> str:
    """
    Embed sentences with all-MiniLM-L6-v2, cluster with AgglomerativeClustering
    (cosine metric, threshold=0.7, no UMAP), find 5 nearest centroids per cluster,
    generate 4 Plotly charts, and save summaries.json + emb.npy.

    Args:
        loaded_data_path: Path to the JSON saved by load_scopus_csv.

    Returns:
        JSON string with cluster stats and chart file paths.
    """
    data      = json.loads(Path(loaded_data_path).read_text())
    sentences = data["sentences"]

    model      = SentenceTransformer(MODEL_NAME)
    embeddings = model.encode(sentences, normalize_embeddings=True, show_progress_bar=False)

    clustering = AgglomerativeClustering(
        n_clusters=None,
        metric="cosine",
        linkage="average",
        distance_threshold=0.7,
    )
    labels = clustering.fit_predict(embeddings)

    unique_labels  = list(set(labels.tolist()))
    n_topics       = len(unique_labels)

    # Centroids: mean of each cluster's embeddings
    centroids = np.array(list(map(
        lambda lbl: embeddings[labels == lbl].mean(axis=0),
        unique_labels,
    )))

    # Top-N nearest sentences to each centroid
    def _top_n_for_cluster(lbl):
        mask     = np.where(labels == lbl)[0]
        c_embs   = embeddings[mask]
        centroid = c_embs.mean(axis=0, keepdims=True)
        sims     = cosine_similarity(centroid, c_embs)[0]
        top_idx  = np.argsort(sims)[::-1][:N_CENTROIDS]
        return list(map(lambda i: sentences[mask[i]], top_idx))

    top_evidence = dict(zip(unique_labels, list(map(_top_n_for_cluster, unique_labels))))

    # Cluster sizes
    sizes = list(map(lambda lbl: int((labels == lbl).sum()), unique_labels))

    # ── Chart 1: Bar — cluster sizes ────────────────────────────────────
    fig1 = px.bar(
        x=list(map(str, unique_labels)),
        y=sizes,
        labels={"x": "Cluster", "y": "Sentences"},
        title="Cluster Size Distribution",
        template="plotly_dark",
        color=sizes,
        color_continuous_scale="Viridis",
    )
    chart1_path = str(ARTIFACTS_DIR / "chart_cluster_sizes.html")
    fig1.write_html(chart1_path, full_html=False)

    # ── Chart 2: Scatter — 2-D PCA projection ───────────────────────────
    from sklearn.decomposition import PCA
    coords = PCA(n_components=2, random_state=42).fit_transform(embeddings)
    fig2 = go.Figure(go.Scatter(
        x=coords[:, 0], y=coords[:, 1],
        mode="markers",
        marker=dict(color=labels.tolist(), colorscale="Turbo", size=4, opacity=0.7),
        text=list(map(lambda i: f"Cluster {labels[i]}", range(len(labels)))),
    ))
    fig2.update_layout(title="Embedding Space (PCA 2D)", template="plotly_dark")
    chart2_path = str(ARTIFACTS_DIR / "chart_pca_scatter.html")
    fig2.write_html(chart2_path, full_html=False)

    # ── Chart 3: Pie — top-10 clusters by size ───────────────────────────
    top10_idx    = np.argsort(sizes)[::-1][:10].tolist()
    top10_labels = list(map(lambda i: f"Cluster {unique_labels[i]}", top10_idx))
    top10_sizes  = list(map(lambda i: sizes[i], top10_idx))
    fig3 = px.pie(names=top10_labels, values=top10_sizes,
                  title="Top 10 Clusters by Size", template="plotly_dark")
    chart3_path = str(ARTIFACTS_DIR / "chart_top10_pie.html")
    fig3.write_html(chart3_path, full_html=False)

    # ── Chart 4: Heatmap — centroid similarity matrix (top 20) ──────────
    top20    = min(20, n_topics)
    sim_mat  = cosine_similarity(centroids[:top20])
    fig4     = px.imshow(
        sim_mat,
        labels=dict(color="Cosine Sim"),
        title=f"Centroid Similarity Heatmap (top {top20})",
        template="plotly_dark",
        color_continuous_scale="RdBu_r",
    )
    chart4_path = str(ARTIFACTS_DIR / "chart_centroid_heatmap.html")
    fig4.write_html(chart4_path, full_html=False)

    # ── Save artefacts ───────────────────────────────────────────────────
    emb_path = str(ARTIFACTS_DIR / "emb.npy")
    np.save(emb_path, embeddings)

    summaries = list(map(lambda lbl: {
        "topic_id":     int(lbl),
        "size":         int((labels == lbl).sum()),
        "top_evidence": top_evidence[lbl],
    }, unique_labels))

    summaries_path = str(ARTIFACTS_DIR / "summaries.json")
    Path(summaries_path).write_text(json.dumps(summaries, ensure_ascii=False, indent=2))

    return json.dumps({
        "n_topics":      n_topics,
        "total_sents":   len(sentences),
        "summaries_path": summaries_path,
        "emb_path":      emb_path,
        "charts": {
            "cluster_sizes":      chart1_path,
            "pca_scatter":        chart2_path,
            "top10_pie":          chart3_path,
            "centroid_heatmap":   chart4_path,
        },
    })


# ──────────────────────────────────────────────────────────────────────────────
# Tool 3 — label_topics_with_llm
# ──────────────────────────────────────────────────────────────────────────────

LABEL_PROMPT = PromptTemplate.from_template(
    """You are a research librarian labelling academic topics.
For each topic below, return a SHORT label (≤ 8 words) and a one-sentence description.

Topics (JSON list, each with topic_id and top_evidence):
{topics_json}

Respond ONLY with a valid JSON array — no markdown fences, no preamble.
Each element must have: topic_id (int), label (str), description (str).
"""
)


@tool
def label_topics_with_llm(summaries_path: str, top_n: int = 100) -> str:
    """
    Send the top-N (default 100) topics to Mistral via PromptTemplate +
    JsonOutputParser to generate concise labels and descriptions.

    Args:
        summaries_path: Path to summaries.json produced by run_bertopic_discovery.
        top_n:          Number of largest topics to label (max 100).

    Returns:
        JSON string with labelled topics and path to saved labels file.
    """
    summaries = json.loads(Path(summaries_path).read_text())
    sorted_s  = sorted(summaries, key=lambda x: x["size"], reverse=True)
    batch     = sorted_s[:min(top_n, 100)]

    chain  = LABEL_PROMPT | _get_llm() | JsonOutputParser()
    result = chain.invoke({"topics_json": json.dumps(batch, ensure_ascii=False)})

    labels_path = str(ARTIFACTS_DIR / "topic_labels.json")
    Path(labels_path).write_text(json.dumps(result, ensure_ascii=False, indent=2))

    return json.dumps({"labelled_count": len(result), "labels_path": labels_path})


# ──────────────────────────────────────────────────────────────────────────────
# Tool 4 — consolidate_into_themes
# ──────────────────────────────────────────────────────────────────────────────

@tool
def consolidate_into_themes(
    labels_path: str,
    summaries_path: str,
    emb_path: str,
    approved_groups: str,
) -> str:
    """
    Merge approved topic groups into themes, recompute centroids.

    Args:
        labels_path:     Path to topic_labels.json.
        summaries_path:  Path to summaries.json.
        emb_path:        Path to emb.npy.
        approved_groups: JSON string — list of groups, each group is a list of
                         topic_ids to merge: e.g. "[[0,3,7],[1,5],[2]]".

    Returns:
        JSON string with theme count, theme details, and saved themes path.
    """
    labels    = json.loads(Path(labels_path).read_text())
    summaries = json.loads(Path(summaries_path).read_text())
    embeddings = np.load(emb_path)

    groups      = json.loads(approved_groups)
    label_map   = {item["topic_id"]: item for item in labels}
    summary_map = {item["topic_id"]: item for item in summaries}

    def _build_theme(idx_group: tuple) -> dict:
        theme_idx, group = idx_group
        member_ids    = group
        member_labels = list(map(lambda tid: label_map.get(tid, {}).get("label", f"Topic {tid}"), member_ids))
        all_evidence  = sum(list(map(lambda tid: summary_map.get(tid, {}).get("top_evidence", []), member_ids)), [])
        total_size    = sum(list(map(lambda tid: summary_map.get(tid, {}).get("size", 0), member_ids)))

        # recompute centroid from member topic centroids
        member_centroids = np.array(list(map(
            lambda tid: embeddings[np.array([], dtype=int)].mean(axis=0)  # placeholder
            if summary_map.get(tid, {}).get("size", 0) == 0
            else np.zeros(embeddings.shape[1]),  # fallback zero vector
            member_ids,
        )))
        centroid = member_centroids.mean(axis=0).tolist()

        return {
            "theme_id":      theme_idx,
            "topic_ids":     member_ids,
            "member_labels": member_labels,
            "theme_label":   member_labels[0],
            "total_size":    total_size,
            "top_evidence":  all_evidence[:N_CENTROIDS],
            "centroid":      centroid,
        }

    themes = list(map(_build_theme, enumerate(groups)))

    themes_path = str(ARTIFACTS_DIR / "themes.json")
    Path(themes_path).write_text(json.dumps(themes, ensure_ascii=False, indent=2))

    return json.dumps({"theme_count": len(themes), "themes_path": themes_path})


# ──────────────────────────────────────────────────────────────────────────────
# Tool 5 — compare_with_taxonomy
# ──────────────────────────────────────────────────────────────────────────────

TAXONOMY_PROMPT = PromptTemplate.from_template(
    """You are a research classifier mapping discovered themes to the PAJAIS taxonomy.

PAJAIS categories:
{categories}

Discovered themes (JSON):
{themes_json}

For each theme, select the SINGLE best-matching PAJAIS category.
Respond ONLY with a valid JSON array — no markdown, no preamble.
Each element: theme_id (int), theme_label (str), pajais_category (str), confidence (0-1 float), rationale (str ≤ 20 words).
"""
)


@tool
def compare_with_taxonomy(themes_path: str) -> str:
    """
    Map consolidated themes to the PAJAIS 25-category taxonomy via Mistral.

    Args:
        themes_path: Path to themes.json produced by consolidate_into_themes.

    Returns:
        JSON string with mapping results and saved taxonomy comparison path.
    """
    themes = json.loads(Path(themes_path).read_text())

    chain  = TAXONOMY_PROMPT | _get_llm() | JsonOutputParser()
    result = chain.invoke({
        "categories":  "\n".join(list(map(lambda c: f"- {c}", PAJAIS_CATEGORIES))),
        "themes_json": json.dumps(themes, ensure_ascii=False),
    })

    taxonomy_path = str(ARTIFACTS_DIR / "taxonomy_mapping.json")
    Path(taxonomy_path).write_text(json.dumps(result, ensure_ascii=False, indent=2))

    return json.dumps({"mapped_count": len(result), "taxonomy_path": taxonomy_path})


# ──────────────────────────────────────────────────────────────────────────────
# Tool 6 — generate_comparison_csv
# ──────────────────────────────────────────────────────────────────────────────

@tool
def generate_comparison_csv(csv_path: str, taxonomy_path: str) -> str:
    """
    Generate a side-by-side comparison CSV of abstract vs title analysis results.

    Args:
        csv_path:      Path to the original Scopus CSV.
        taxonomy_path: Path to taxonomy_mapping.json for the abstract run.

    Returns:
        JSON string with row count and path to the comparison CSV.
    """
    df      = pd.read_csv(csv_path)
    mapping = json.loads(Path(taxonomy_path).read_text())

    abstract_col = next(filter(lambda c: c in df.columns, RUN_CONFIGS["abstract"]), None)
    title_col    = next(filter(lambda c: c in df.columns, RUN_CONFIGS["title"]), None)

    abstracts = list(map(lambda t: _clean_text(str(t)), df[abstract_col].fillna("").tolist()))
    titles    = list(map(lambda t: _clean_text(str(t)), df[title_col].fillna("").tolist()))

    category_labels = list(map(lambda m: m.get("pajais_category", "Unclassified"), mapping))
    padded_cats     = (category_labels + ["Unclassified"] * len(df))[:len(df)]

    comparison_df = pd.DataFrame({
        "paper_id":         list(range(1, len(df) + 1)),
        "title":            titles,
        "abstract_snippet": list(map(lambda a: a[:200], abstracts)),
        "pajais_category":  padded_cats,
        "confidence":       list(map(lambda m: m.get("confidence", 0.0), mapping))[:len(df)]
                            + [0.0] * max(0, len(df) - len(mapping)),
    })

    out_path = str(ARTIFACTS_DIR / "abstract_vs_title_comparison.csv")
    comparison_df.to_csv(out_path, index=False)

    return json.dumps({"rows": len(comparison_df), "comparison_csv": out_path})


# ──────────────────────────────────────────────────────────────────────────────
# Tool 7 — export_narrative
# ──────────────────────────────────────────────────────────────────────────────

NARRATIVE_PROMPT = PromptTemplate.from_template(
    """You are an academic author writing Section 7 (Discussion & Implications) of a
systematic literature review on AI in business and management journals.

Use the taxonomy mapping below as your evidence base.

Taxonomy mapping (JSON):
{taxonomy_json}

Write exactly ~500 words as flowing academic prose (no bullet points, no headers).
Discuss: (1) dominant themes, (2) gaps relative to the PAJAIS taxonomy,
(3) methodological implications, (4) future research directions.
Cite themes by their label. Maintain formal academic register throughout.
"""
)


@tool
def export_narrative(taxonomy_path: str) -> str:
    """
    Generate a ~500-word Section 7 narrative via Mistral and save it as a text file.

    Args:
        taxonomy_path: Path to taxonomy_mapping.json.

    Returns:
        JSON string with word count and path to the saved narrative file.
    """
    taxonomy = json.loads(Path(taxonomy_path).read_text())

    chain    = NARRATIVE_PROMPT | _get_llm()
    response = chain.invoke({"taxonomy_json": json.dumps(taxonomy, ensure_ascii=False)})

    narrative_text = response.content

    narrative_path = str(ARTIFACTS_DIR / "section7_narrative.txt")
    Path(narrative_path).write_text(narrative_text, encoding="utf-8")

    word_count = len(narrative_text.split())

    return json.dumps({
        "word_count":     word_count,
        "narrative_path": narrative_path,
        "preview":        narrative_text[:300] + "…",
    })