"""
research_chatbot.py — Research Intelligence Chatbot engine.

Builds a compressed context summary from all pipeline outputs, then
uses the Groq LLM to answer research questions in an academic style.
Caches the context summary so the LLM prompt stays small and fast.

Includes:
  - Deep optimization analysis (cluster density, balance, overlap)
  - Keyword redundancy detection
  - Silhouette score estimation
  - Outlier identification
  - Parameter recommendation logic
  - Enriched context from all pipeline outputs
"""

import os
import json
import logging
import hashlib
from typing import Dict, Any, List, Optional
from collections import Counter

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# --------------------------------------------------------------------------
# Suggested prompts — visible in the sidebar
# --------------------------------------------------------------------------
SUGGESTED_PROMPTS = [
    "What is the optimal number of clusters for this dataset?",
    "Which computational strategy dominates across papers?",
    "Which clusters overlap semantically and should be merged?",
    "What research gaps exist in this corpus?",
    "Which papers are potential outliers?",
    "How can clustering quality be improved?",
    "Summarize the largest cluster in academic language.",
    "What HDBSCAN parameters would produce better results?",
    "Which topics are novel vs well-established?",
    "What trends do you see across strategy types?",
    "Identify weak or poorly formed clusters.",
    "Suggest better UMAP settings for this data.",
    "What preprocessing improvements could help?",
    "Which clusters have the highest keyword redundancy?",
    "Analyze the cluster balance and distribution.",
    "What are the dominant research themes?",
]


# --------------------------------------------------------------------------
# Optimization analysis helpers
# --------------------------------------------------------------------------

def _compute_cluster_balance(labels: List[int]) -> Dict[str, Any]:
    """Analyze how balanced the cluster sizes are."""
    if not labels:
        return {}
    counts = Counter(labels)
    # Remove noise
    counts.pop(-1, None)
    if not counts:
        return {}

    sizes = list(counts.values())
    mean_sz = np.mean(sizes)
    std_sz = np.std(sizes)
    min_sz = min(sizes)
    max_sz = max(sizes)
    cv = std_sz / mean_sz if mean_sz > 0 else 0  # coefficient of variation

    # Identify imbalanced clusters
    tiny = [cid for cid, sz in counts.items() if sz < mean_sz * 0.3]
    large = [cid for cid, sz in counts.items() if sz > mean_sz * 2.5]

    return {
        "mean_size": round(mean_sz, 1),
        "std_size": round(std_sz, 1),
        "min_size": min_sz,
        "max_size": max_sz,
        "cv": round(cv, 3),
        "balance_quality": (
            "Well-balanced" if cv < 0.5
            else "Moderately imbalanced" if cv < 1.0
            else "Highly imbalanced"
        ),
        "tiny_clusters": tiny,
        "oversized_clusters": large,
    }


def _detect_keyword_redundancy(cluster_details: Dict) -> List[Dict]:
    """Find clusters with high keyword overlap indicating potential merges."""
    overlaps = []
    cluster_ids = sorted(cluster_details.keys())
    for i, cid_a in enumerate(cluster_ids):
        kw_a = set(k.strip().lower() for k in cluster_details[cid_a].get("keywords", "").split(","))
        for cid_b in cluster_ids[i + 1:]:
            kw_b = set(k.strip().lower() for k in cluster_details[cid_b].get("keywords", "").split(","))
            intersection = kw_a & kw_b
            union = kw_a | kw_b
            if union:
                jaccard = len(intersection) / len(union)
                if jaccard >= 0.3:
                    overlaps.append({
                        "cluster_a": cid_a,
                        "cluster_b": cid_b,
                        "jaccard_sim": round(jaccard, 3),
                        "shared_keywords": list(intersection)[:8],
                    })
    return sorted(overlaps, key=lambda x: x["jaccard_sim"], reverse=True)


def _estimate_silhouette_info(
    embeddings_2d: Optional[List] = None,
    cluster_labels: Optional[List] = None,
) -> Dict[str, Any]:
    """Estimate silhouette-like quality metrics from 2D embeddings."""
    if not embeddings_2d or not cluster_labels:
        return {}

    try:
        from sklearn.metrics import silhouette_score, calinski_harabasz_score

        arr = np.array(embeddings_2d)
        labels_arr = np.array(cluster_labels)

        # Filter out noise
        valid = labels_arr >= 0
        if valid.sum() < 10 or len(set(labels_arr[valid])) < 2:
            return {"note": "Too few valid clusters for quality scoring"}

        sil = silhouette_score(arr[valid], labels_arr[valid])
        ch = calinski_harabasz_score(arr[valid], labels_arr[valid])

        quality = (
            "Excellent" if sil > 0.5
            else "Good" if sil > 0.3
            else "Fair" if sil > 0.1
            else "Poor"
        )

        return {
            "silhouette_score_2d": round(sil, 4),
            "calinski_harabasz": round(ch, 2),
            "quality_rating": quality,
            "interpretation": (
                f"Silhouette = {sil:.3f} ({quality}). "
                f"{'Clusters are well-separated.' if sil > 0.3 else 'Some clusters overlap significantly.'} "
                f"Calinski-Harabasz = {ch:.1f} (higher = more distinct clusters)."
            ),
        }
    except Exception as e:
        logger.warning("Could not compute silhouette: %s", e)
        return {"error": str(e)}


def _identify_outlier_clusters(cluster_details: Dict, balance: Dict) -> List[Dict]:
    """Identify clusters that may be outlier groups."""
    outliers = []
    if not cluster_details or not balance:
        return outliers

    mean_sz = balance.get("mean_size", 0)
    for cid, detail in cluster_details.items():
        n_docs = detail.get("n_documents", 0)
        reasons = []

        if n_docs <= 2:
            reasons.append("Only 1–2 documents (singleton/pair)")
        elif mean_sz > 0 and n_docs < mean_sz * 0.15:
            reasons.append(f"Very small ({n_docs} docs vs mean {mean_sz:.0f})")

        kws = detail.get("keywords", "")
        if len(kws.split(",")) <= 2:
            reasons.append("Very few distinguishing keywords")

        if reasons:
            outliers.append({
                "cluster_id": cid,
                "n_documents": n_docs,
                "reasons": reasons,
            })

    return outliers


def _generate_parameter_recommendations(
    n_clusters: int,
    n_papers: int,
    hdbscan_params: Dict,
    quality: Dict,
    balance: Dict,
    overlaps: List,
) -> List[str]:
    """Generate actionable parameter recommendations."""
    recs = []

    sil = quality.get("silhouette_score_2d", None)

    # Cluster count vs paper count ratio
    if n_papers > 0 and n_clusters > 0:
        ratio = n_papers / n_clusters
        if ratio < 5:
            recs.append(
                f"📉 High cluster count ({n_clusters}) relative to papers ({n_papers}). "
                f"Increase HDBSCAN `min_cluster_size` to 10–15 to reduce fragmentation."
            )
        elif ratio > 100:
            recs.append(
                f"📈 Low cluster count ({n_clusters}) for {n_papers} papers. "
                f"Decrease `min_cluster_size` to 3–5 to capture more nuanced groupings."
            )

    # Silhouette guidance
    if sil is not None:
        if sil < 0.1:
            recs.append(
                "⚠️ Very low silhouette score — clusters overlap heavily. "
                "Try: (1) increase UMAP `n_neighbors` to 20–30, "
                "(2) use `min_dist=0.0` in UMAP, "
                "(3) increase `min_cluster_size`."
            )
        elif sil < 0.3:
            recs.append(
                "🔄 Moderate silhouette — some overlap remains. "
                "Consider merging overlapping clusters or adjusting `min_samples`."
            )

    # Balance
    cv = balance.get("cv", 0)
    if cv > 1.0:
        recs.append(
            f"⚖️ High size imbalance (CV={cv:.2f}). "
            f"HDBSCAN `cluster_selection_method='leaf'` may produce more even clusters."
        )

    # Overlaps
    if len(overlaps) >= 3:
        recs.append(
            f"🔗 {len(overlaps)} keyword-overlapping cluster pairs detected. "
            f"Lower the merge threshold (currently 0.75) to 0.60–0.65."
        )

    # HDBSCAN params
    mcs = hdbscan_params.get("min_cluster_size", "?")
    ms = hdbscan_params.get("min_samples", "?")
    recs.append(
        f"Current HDBSCAN: min_cluster_size={mcs}, min_samples={ms}. "
        f"Typical range for academic corpora: mcs=5–15, ms=2–5."
    )

    if not recs:
        recs.append("✅ Current parameters appear reasonable. No major issues detected.")

    return recs


# --------------------------------------------------------------------------
# Context builder — distills all pipeline outputs into a compact summary
# --------------------------------------------------------------------------

def build_context_summary(
    result=None,
    comp_data: Optional[Dict[str, Any]] = None,
) -> str:
    """
    Build a compact text summary of all available analysis outputs.
    This is injected into the LLM system prompt so it can answer
    questions without sending raw data every time.

    Includes optimization metrics for intelligent recommendations.
    """
    parts: List[str] = []

    # ── Topic-modeling pipeline results ──
    if result is not None and hasattr(result, "status") and result.status == "success":
        ct = result.combined_topics
        parts.append("## Topic Modeling Pipeline (SPECTER2 + HDBSCAN)")
        parts.append(f"- Status: {result.status}")
        parts.append(f"- Pipeline mode: {getattr(result, 'pipeline_mode', 'default')}")

        has_cluster = "cluster_id" in ct.columns
        n_clusters = 0
        if has_cluster:
            n_clusters = int(ct["cluster_id"].nunique())
            parts.append(f"- Clusters: {n_clusters}")
        parts.append(f"- Total topics/sub-topics: {len(ct)}")
        parts.append(f"- Total documents in dataset: {ct['n_documents'].sum() if 'n_documents' in ct.columns else '?'}")

        # Cluster details
        cd = getattr(result, "cluster_data", {})
        cluster_details = cd.get("cluster_details", {})
        if cluster_details:
            parts.append("\n### Cluster Details")
            for cid in sorted(cluster_details.keys()):
                d = cluster_details[cid]
                parts.append(
                    f"  Cluster {cid}: {d.get('n_documents', '?')} docs | "
                    f"Keywords: {d.get('keywords', '')[:100]}"
                )
                reps = d.get("representatives", [])
                if reps:
                    parts.append(f"    Representative: {reps[0][:120]}")

        # HDBSCAN params
        hparams = cd.get("hdbscan_params", {})
        if hparams:
            parts.append(f"\n### HDBSCAN Parameters: {hparams}")
        n_clust = cd.get("n_clusters")
        if n_clust:
            parts.append(f"- Final cluster count: {n_clust}")

        # ── Optimization Metrics ──
        emb_2d = cd.get("embeddings_2d")
        cl_labels = cd.get("cluster_labels")

        # Cluster balance
        if cl_labels:
            balance = _compute_cluster_balance(cl_labels)
            if balance:
                parts.append("\n### Cluster Balance Analysis")
                parts.append(f"  Mean size: {balance['mean_size']} | Std: {balance['std_size']} | CV: {balance['cv']}")
                parts.append(f"  Min: {balance['min_size']} | Max: {balance['max_size']}")
                parts.append(f"  Balance quality: {balance['balance_quality']}")
                if balance.get("tiny_clusters"):
                    parts.append(f"  ⚠ Tiny clusters (< 30% mean): {balance['tiny_clusters']}")
                if balance.get("oversized_clusters"):
                    parts.append(f"  ⚠ Oversized clusters (> 2.5× mean): {balance['oversized_clusters']}")

        # Silhouette & quality
        quality = _estimate_silhouette_info(emb_2d, cl_labels)
        if quality and "silhouette_score_2d" in quality:
            parts.append("\n### Clustering Quality Metrics")
            parts.append(f"  Silhouette Score (2D): {quality['silhouette_score_2d']}")
            parts.append(f"  Calinski-Harabasz: {quality['calinski_harabasz']}")
            parts.append(f"  Quality Rating: {quality['quality_rating']}")
            parts.append(f"  {quality['interpretation']}")

        # Keyword redundancy / overlap
        if cluster_details:
            overlaps = _detect_keyword_redundancy(cluster_details)
            if overlaps:
                parts.append(f"\n### Keyword Overlap Analysis ({len(overlaps)} overlapping pairs)")
                for ov in overlaps[:5]:
                    parts.append(
                        f"  C{ov['cluster_a']} ↔ C{ov['cluster_b']}: "
                        f"Jaccard={ov['jaccard_sim']} | Shared: {', '.join(ov['shared_keywords'][:5])}"
                    )
            else:
                parts.append("\n### Keyword Overlap: No significant overlap detected between clusters.")

            # Outlier clusters
            if cl_labels:
                balance = _compute_cluster_balance(cl_labels)
                outlier_clusters = _identify_outlier_clusters(cluster_details, balance)
                if outlier_clusters:
                    parts.append(f"\n### Potential Outlier Clusters ({len(outlier_clusters)})")
                    for oc in outlier_clusters:
                        parts.append(
                            f"  Cluster {oc['cluster_id']}: {oc['n_documents']} docs — "
                            f"{'; '.join(oc['reasons'])}"
                        )

            # Parameter recommendations
            if cl_labels and hparams:
                n_papers = len(cl_labels) if cl_labels else 0
                balance_for_recs = _compute_cluster_balance(cl_labels) if cl_labels else {}
                overlaps_for_recs = _detect_keyword_redundancy(cluster_details)
                recs = _generate_parameter_recommendations(
                    n_clusters, n_papers, hparams, quality,
                    balance_for_recs, overlaps_for_recs,
                )
                parts.append("\n### Automated Parameter Recommendations")
                for rec in recs:
                    parts.append(f"  {rec}")

        # Taxonomy summary
        tax = getattr(result, "taxonomy_map", {})
        meta = tax.get("metadata", {})
        if meta:
            parts.append(f"\n### Taxonomy: {meta.get('mapped_count', 0)} mapped, "
                         f"{meta.get('novel_count', 0)} novel (threshold {meta.get('threshold', 0.15)})")

        # Topic keywords sample (top 15)
        if "keywords" in ct.columns:
            parts.append("\n### Top 15 Topic Keywords")
            for _, row in ct.head(15).iterrows():
                lbl = row.get("label", row.get("cluster_label", ""))
                parts.append(f"  Topic {row.get('topic_id', '?')}: {row['keywords'][:100]} [{lbl}]")

        # Graph edges summary
        edges = cd.get("graph_edges", [])
        if edges:
            top = sorted(edges, key=lambda e: e["weight"], reverse=True)[:8]
            parts.append("\n### Top Inter-cluster Similarities")
            for e in top:
                parts.append(f"  C{e['source']} ↔ C{e['target']}: {e['weight']:.3f}")

        # Council debates summary (if available)
        debates = getattr(result, "council_debates", {})
        if debates:
            parts.append(f"\n### AI Council Labels ({len(debates)} clusters)")
            for cid in sorted(debates.keys()):
                d = debates[cid]
                parts.append(
                    f"  C{cid}: Expert='{d.get('expert', '?')}' | "
                    f"Analyst='{d.get('analyst', '?')}' | "
                    f"Final='{d.get('final_label', '?')}'"
                )

    # ── Computation Strategy clustering ──
    if comp_data is not None:
        parts.append("\n## Computational Strategy Clustering")
        parts.append(f"- Papers: {comp_data.get('n_papers', '?')}")
        parts.append(f"- Clusters: {comp_data.get('n_clusters', '?')}")
        parts.append(f"- Strategy types: {', '.join(comp_data.get('strategy_types', []))}")

        cn = comp_data.get("cluster_names", {})
        ck = comp_data.get("cluster_keywords", {})
        df = comp_data.get("df")

        # Strategy type distribution
        if df is not None and "Strategy Type" in df.columns:
            type_counts = df["Strategy Type"].value_counts().to_dict()
            parts.append("\n### Strategy Type Distribution")
            for stype, count in sorted(type_counts.items(), key=lambda x: x[1], reverse=True):
                if stype.strip():
                    parts.append(f"  {stype}: {count} papers")

        for cid in sorted(cn.keys()):
            n = int((df["cluster_id"] == cid).sum()) if df is not None else "?"
            kws = ", ".join(ck.get(cid, [])[:6])
            parts.append(f"  Cluster {cid} ({cn[cid]}): {n} papers | {kws}")

    if not parts:
        return "No analysis has been completed yet. Run the pipeline first."

    return "\n".join(parts)


# --------------------------------------------------------------------------
# Chat with Groq
# --------------------------------------------------------------------------

SYSTEM_PROMPT = """You are the **Research Intelligence Assistant** embedded in a PCCM
(Paper Computational Clustering & Modeling) platform. You are an expert in:
- Interpreting SPECTER2 + HDBSCAN clustering results
- Research methodology, gap analysis, and trend identification
- Dimensionality reduction (UMAP) and cluster validation metrics
- Academic publishing patterns and computational strategy analysis

## Your Behavioral Rules

1. **Ground every answer in the provided analysis context** — never hallucinate data.
   If the context doesn't contain enough information, explicitly say so.

2. **Respond in concise, structured academic language** suitable for IS/CS researchers.
   Use bullet points, numbered lists, and short paragraphs.

3. **When suggesting parameter changes**, always explain:
   - What the current setting is
   - What you recommend changing it to
   - The expected effect on cluster quality
   - Any trade-offs involved

4. **For cluster interpretation**, always reference:
   - Cluster ID, keyword composition, document count
   - Semantic overlap with neighboring clusters
   - Whether the cluster represents a novel or established theme

5. **For optimization questions**, leverage the provided metrics:
   - Silhouette score and its meaning
   - Cluster balance (CV, tiny/oversized clusters)
   - Keyword overlap (Jaccard) between clusters
   - Parameter recommendations from the automated analysis

6. **Keep responses under 500 words** unless the user asks for detailed analysis.

7. **Use markdown formatting** for readability:
   - Bold for emphasis
   - Code blocks for parameter values
   - Tables where comparing multiple options

8. **When identifying research gaps**, consider:
   - Under-represented clusters (few documents)
   - Novel themes from taxonomy classification
   - Strategy types with low coverage
   - Keywords that appear in few topics

=== CURRENT ANALYSIS CONTEXT ===
{context}
"""


def chat_with_groq(
    messages: List[Dict[str, str]],
    context_summary: str,
    groq_api_key: str,
    model: str = "llama-3.3-70b-versatile",
) -> str:
    """
    Send a conversation to Groq with the analysis context injected.
    Returns the assistant reply text.
    """
    try:
        from groq import Groq
    except ImportError:
        return "❌ `groq` package not installed. Add it to requirements.txt."

    if not groq_api_key:
        return "❌ No Groq API key provided. Enter it in the sidebar."

    client = Groq(api_key=groq_api_key)

    system_msg = {
        "role": "system",
        "content": SYSTEM_PROMPT.format(context=context_summary),
    }

    # Build messages list: system + conversation history (last 20 messages = ~10 turns)
    trimmed = messages[-20:]
    full_messages = [system_msg] + trimmed

    try:
        resp = client.chat.completions.create(
            model=model,
            messages=full_messages,
            temperature=0.35,
            max_tokens=1500,
        )
        return resp.choices[0].message.content.strip()
    except Exception as e:
        logger.error("Groq chat error: %s", e)
        return f"❌ LLM error: {e}"


def get_context_stats(result=None, comp_data=None) -> Dict[str, Any]:
    """
    Return a dict of quick stats for the right panel display.
    This avoids re-computing in the UI layer.
    """
    stats = {
        "n_clusters": 0,
        "n_topics": 0,
        "n_mapped": 0,
        "n_novel": 0,
        "silhouette": None,
        "quality_rating": None,
        "balance_quality": None,
        "n_overlaps": 0,
        "n_outliers": 0,
        "comp_papers": 0,
        "comp_clusters": 0,
        "comp_strategy_types": 0,
    }

    if result is not None and hasattr(result, "status") and result.status == "success":
        ct = result.combined_topics
        has_cluster = "cluster_id" in ct.columns
        stats["n_clusters"] = int(ct["cluster_id"].nunique()) if has_cluster else 0
        stats["n_topics"] = len(ct)

        tax_m = getattr(result, "taxonomy_map", {}).get("metadata", {})
        stats["n_mapped"] = tax_m.get("mapped_count", 0)
        stats["n_novel"] = tax_m.get("novel_count", 0)

        cd = getattr(result, "cluster_data", {})
        emb_2d = cd.get("embeddings_2d")
        cl_labels = cd.get("cluster_labels")
        cluster_details = cd.get("cluster_details", {})

        if emb_2d and cl_labels:
            quality = _estimate_silhouette_info(emb_2d, cl_labels)
            stats["silhouette"] = quality.get("silhouette_score_2d")
            stats["quality_rating"] = quality.get("quality_rating")

            balance = _compute_cluster_balance(cl_labels)
            stats["balance_quality"] = balance.get("balance_quality")

        if cluster_details:
            overlaps = _detect_keyword_redundancy(cluster_details)
            stats["n_overlaps"] = len(overlaps)

            if cl_labels:
                balance = _compute_cluster_balance(cl_labels)
                outliers = _identify_outlier_clusters(cluster_details, balance)
                stats["n_outliers"] = len(outliers)

    if comp_data is not None:
        stats["comp_papers"] = comp_data.get("n_papers", 0)
        stats["comp_clusters"] = comp_data.get("n_clusters", 0)
        stats["comp_strategy_types"] = len(comp_data.get("strategy_types", []))

    return stats