"""
tools.py - BERTopic Agentic AI Tools
Stateless, functional tools with zero try/except blocks.
All errors propagate to LLM via @tool.
"""

import pandas as pd
import numpy as np
import re
import json
from pathlib import Path
from typing import Dict, Any, List

import nltk
from nltk.tokenize import sent_tokenize
from sentence_transformers import SentenceTransformer
from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics.pairwise import cosine_similarity
from bertopic import BERTopic
from bertopic.vectorizers import ClassTfidfTransformer
from bertopic.representation import KeyBERTInspired
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.express as px

from langchain_core.tools import tool
from langchain_mistralai import ChatMistralAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser

# Download NLTK data
nltk.download('punkt', quiet=True)
nltk.download('punkt_tab', quiet=True)

# Global paths
DATA_DIR = Path("./data")
DATA_DIR.mkdir(exist_ok=True)

def clean_abstract_text(text: str) -> str:
    """Remove publisher boilerplate noise from abstracts."""
    patterns = [
        r'©\s*\d{4}.*?rights reserved',
        r'All rights reserved',
        r'Elsevier.*?Ltd',
        r'Springer.*?Verlag',
        r'Published by.*?$',
        r'\[ABSTRACT FROM AUTHOR\]',
        r'\[ABSTRACT FROM PUBLISHER\]',
    ]
    cleaned = text
    cleaned = re.sub('|'.join(patterns), '', cleaned, flags=re.IGNORECASE)
    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    return cleaned

@tool
def load_scopus_csv(filepath: str) -> str:
    """
    Load Scopus CSV, count papers, split abstracts into sentences, remove boilerplate.
    Returns stats string with paper count and sentence count.
    """
    # FIX: Use encoding_errors='replace' to prevent Scopus invisible character crashes
    df = pd.read_csv(filepath, encoding='utf-8-sig', encoding_errors='replace')
    
    paper_count = len(df)
    
    df['Abstract'] = df['Abstract'].fillna('').astype(str)
    df['cleaned_abstract'] = df['Abstract'].apply(clean_abstract_text)
    
    sentences = []
    paper_ids = []
    
    for idx, row in df.iterrows():
        sents = sent_tokenize(row['cleaned_abstract'])
        sentences.extend(sents)
        paper_ids.extend([idx] * len(sents))
        
    sentence_df = pd.DataFrame({
        'sentence': sentences,
        'paper_id': paper_ids
    })
    
    sentence_df.to_csv(DATA_DIR / 'sentences.csv', index=False)
    df.to_csv(DATA_DIR / 'papers.csv', index=False)
    
    return f"Loaded {paper_count} papers, extracted {len(sentences)} sentences. Saved to data/sentences.csv and data/papers.csv"

@tool
def run_bertopic_discovery(run_key: str, threshold: float = 0.7) -> str:
    """
    Embed sentences with all-MiniLM-L6-v2, cluster in 384d space using AgglomerativeClustering.
    NO UMAP dimensionality reduction. Generate Plotly charts and save outputs.
    """
    sentences_df = pd.read_csv(DATA_DIR / 'sentences.csv')
    sentences = sentences_df['sentence'].tolist()
    
    # Embed sentences
    embedder = SentenceTransformer('all-MiniLM-L6-v2')
    embeddings = embedder.encode(sentences, normalize_embeddings=True, show_progress_bar=True)
    
    # Cluster directly in 384d space
    clusterer = AgglomerativeClustering(
        n_clusters=None,
        metric='cosine',
        linkage='average',
        distance_threshold=threshold
    )
    cluster_labels = clusterer.fit_predict(embeddings)
    
    # Build topic summaries
    unique_labels = sorted(set(cluster_labels))
    topic_summaries = []
    
    for label in unique_labels:
        cluster_mask = cluster_labels == label
        cluster_embeddings = embeddings[cluster_mask]
        cluster_sentences_idx = np.where(cluster_mask)[0]
        
        centroid = cluster_embeddings.mean(axis=0).reshape(1, -1)
        similarities = cosine_similarity(cluster_embeddings, centroid).flatten()
        
        top_5_idx = similarities.argsort()[-5:][::-1]
        top_sentences = [sentences[cluster_sentences_idx[i]] for i in top_5_idx]
        
        paper_ids = sentences_df.iloc[cluster_sentences_idx]['paper_id'].unique()
        
        topic_summaries.append({
            'topic_id': int(label),
            'sentence_count': int(cluster_mask.sum()),
            'paper_count': int(len(paper_ids)),
            'top_sentences': top_sentences,
            'centroid': centroid.tolist()
        })
        
    # Save outputs
    run_dir = DATA_DIR / run_key
    run_dir.mkdir(parents=True, exist_ok=True)
    
    with open(run_dir / 'summaries.json', 'w') as f:
        json.dump(topic_summaries, f, indent=2)
        
    np.save(run_dir / 'emb.npy', embeddings)
    np.save(run_dir / 'labels.npy', cluster_labels)
    
    # Generate Plotly charts
    embeddings_2d = embeddings[:, :2]  # Use first 2 dimensions for visualization
    
    # Chart 1: Intertopic Distance Map
    fig1 = px.scatter(
        x=embeddings_2d[:, 0],
        y=embeddings_2d[:, 1],
        color=cluster_labels,
        title='Intertopic Distance Map',
        labels={'x': 'Dimension 1', 'y': 'Dimension 2', 'color': 'Topic'}
    )
    fig1.write_html(run_dir / 'intertopic_map.html', include_plotlyjs='cdn')
    
    # Chart 2: Topic Bars
    topic_sizes = pd.Series(cluster_labels).value_counts().sort_index()
    fig2 = go.Figure(data=[go.Bar(x=topic_sizes.index, y=topic_sizes.values)])
    fig2.update_layout(title='Topic Sizes', xaxis_title='Topic ID', yaxis_title='Sentence Count')
    fig2.write_html(run_dir / 'topic_bars.html', include_plotlyjs='cdn')
    
    # Chart 3: Hierarchy (dendrogram simulation)
    fig3 = go.Figure(data=[go.Scatter(
        x=list(range(len(unique_labels))),
        y=[len([c for c in cluster_labels if c == l]) for l in unique_labels],
        mode='markers+lines',
        marker=dict(size=10)
    )])
    fig3.update_layout(title='Topic Hierarchy', xaxis_title='Topic ID', yaxis_title='Size')
    fig3.write_html(run_dir / 'hierarchy.html', include_plotlyjs='cdn')
    
    # Chart 4: Heatmap (topic similarity)
    topic_centroids = np.array([np.mean(embeddings[cluster_labels == l], axis=0) for l in unique_labels])
    similarity_matrix = cosine_similarity(topic_centroids)
    fig4 = go.Figure(data=go.Heatmap(z=similarity_matrix, x=unique_labels, y=unique_labels))
    fig4.update_layout(title='Topic Similarity Heatmap')
    fig4.write_html(run_dir / 'heatmap.html', include_plotlyjs='cdn')
    
    return f"Discovered {len(unique_labels)} topics. Saved summaries, embeddings, and 4 charts to {run_dir}"

@tool
def label_topics_with_llm(run_key: str) -> str:
    """
    Send top 100 topics to Mistral API for labeling.
    Parse JSON output: label, category, confidence, reasoning, niche.
    """
    run_dir = DATA_DIR / run_key
    
    with open(run_dir / 'summaries.json', 'r') as f:
        summaries = json.load(f)
        
    top_100 = sorted(summaries, key=lambda x: x['sentence_count'], reverse=True)[:100]
    
    llm = ChatMistralAI(model='mistral-small-latest', temperature=0)
    parser = JsonOutputParser()
    
    prompt = PromptTemplate(
        template="""You are a research domain expert. Analyze these topic clusters from academic papers and provide structured labels.
For each topic, return JSON with:
- label: Short research area name (2-5 words)
- category: Broader domain (e.g., "Machine Learning", "Healthcare", "Finance")
- confidence: 0-1 score
- reasoning: Brief explanation
- niche: true if highly specialized, false if mainstream
Topics:
{topics}
Return a JSON array with one object per topic. {format_instructions}""",
        input_variables=['topics'],
        partial_variables={'format_instructions': parser.get_format_instructions()}
    )
    
    topics_text = json.dumps([{
        'topic_id': t['topic_id'],
        'top_sentences': t['top_sentences'][:3]
    } for t in top_100], indent=2)
    
    chain = prompt | llm | parser
    labels = chain.invoke({'topics': topics_text})
    
    with open(run_dir / 'labels.json', 'w') as f:
        json.dump(labels, f, indent=2)
        
    return f"Labeled {len(labels)} topics. Saved to {run_dir / 'labels.json'}"

@tool
def consolidate_into_themes(run_key: str, theme_map: Dict[str, List[int]]) -> str:
    """
    Group topics into researcher-approved themes.
    Recompute centroids, recount sentences/papers.
    """
    run_dir = DATA_DIR / run_key
    
    with open(run_dir / 'summaries.json', 'r') as f:
        summaries = json.load(f)
        
    embeddings = np.load(run_dir / 'emb.npy')
    cluster_labels = np.load(run_dir / 'labels.npy')
    
    summaries_dict = {s['topic_id']: s for s in summaries}
    themes = []
    
    for theme_name, topic_ids in theme_map.items():
        topic_ids = [int(tid) for tid in topic_ids]
        theme_mask = np.isin(cluster_labels, topic_ids)
        theme_embeddings = embeddings[theme_mask]
        
        centroid = theme_embeddings.mean(axis=0)
        sentence_count = int(theme_mask.sum())
        
        paper_ids = set()
        for tid in topic_ids:
            sentences_df = pd.read_csv(DATA_DIR / 'sentences.csv')
            topic_sentences = sentences_df[cluster_labels == tid]
            paper_ids.update(topic_sentences['paper_id'].unique())
            
        themes.append({
            'theme_name': theme_name,
            'topic_ids': topic_ids,
            'sentence_count': sentence_count,
            'paper_count': len(paper_ids),
            'centroid': centroid.tolist()
        })
        
    with open(run_dir / 'themes.json', 'w') as f:
        json.dump(themes, f, indent=2)
        
    return f"Consolidated {len(themes)} themes. Saved to {run_dir / 'themes.json'}"

@tool
def compare_with_taxonomy(run_key: str) -> str:
    """
    Map themes to PAJAIS 25-category taxonomy using Mistral.
    Return pajais_match, confidence, reasoning, is_novel.
    """
    run_dir = DATA_DIR / run_key
    
    with open(run_dir / 'themes.json', 'r') as f:
        themes = json.load(f)
        
    pajais_categories = [
        "Artificial Intelligence", "Blockchain", "Business Analytics", "Cloud Computing",
        "Cybersecurity", "Data Science", "Digital Innovation", "E-Commerce", "Enterprise Systems",
        "Ethics and Privacy", "Healthcare IT", "Human-Computer Interaction", "Internet of Things",
        "IT Governance", "Knowledge Management", "Mobile Computing", "Social Media",
        "Software Development", "Strategic IS", "Supply Chain", "Technology Adoption",
        "Telecommunications", "Virtual Reality", "Web Technologies", "NOVEL"
    ]
    
    llm = ChatMistralAI(model='mistral-small-latest', temperature=0)
    parser = JsonOutputParser()
    
    prompt = PromptTemplate(
        template="""Map each research theme to the closest PAJAIS category.
PAJAIS Categories: {categories}
Themes:
{themes}
For each theme, return JSON with:
- theme_name: Original theme name
- pajais_match: Best matching category (or NOVEL if none fit)
- confidence: 0-1 score
- reasoning: Explanation
- is_novel: true if NOVEL, false otherwise
Return a JSON array. {format_instructions}""",
        input_variables=['themes', 'categories'],
        partial_variables={'format_instructions': parser.get_format_instructions()}
    )
    
    themes_text = json.dumps([{
        'theme_name': t['theme_name'],
        'sentence_count': t['sentence_count'],
        'paper_count': t['paper_count']
    } for t in themes], indent=2)
    
    chain = prompt | llm | parser
    mappings = chain.invoke({
        'themes': themes_text,
        'categories': ', '.join(pajais_categories)
    })
    
    with open(run_dir / 'taxonomy_map.json', 'w') as f:
        json.dump(mappings, f, indent=2)
        
    return f"Mapped {len(mappings)} themes to taxonomy. Saved to {run_dir / 'taxonomy_map.json'}"

@tool
def generate_comparison_csv() -> str:
    """
    Load themes from abstract and title runs.
    Create side-by-side comparison DataFrame.
    """
    abstract_dir = DATA_DIR / 'abstract_run'
    title_dir = DATA_DIR / 'title_run'
    
    with open(abstract_dir / 'themes.json', 'r') as f:
        abstract_themes = json.load(f)
        
    with open(title_dir / 'themes.json', 'r') as f:
        title_themes = json.load(f)
        
    rows = []
    max_len = max(len(abstract_themes), len(title_themes))
    
    for i in range(max_len):
        row = {}
        
        if i < len(abstract_themes):
            at = abstract_themes[i]
            row['abstract_theme'] = at['theme_name']
            row['abstract_sentences'] = at['sentence_count']
            row['abstract_papers'] = at['paper_count']
        else:
            row['abstract_theme'] = ''
            row['abstract_sentences'] = 0
            row['abstract_papers'] = 0
            
        if i < len(title_themes):
            tt = title_themes[i]
            row['title_theme'] = tt['theme_name']
            row['title_sentences'] = tt['sentence_count']
            row['title_papers'] = tt['paper_count']
        else:
            row['title_theme'] = ''
            row['title_sentences'] = 0
            row['title_papers'] = 0
            
        rows.append(row)
        
    comparison_df = pd.DataFrame(rows)
    comparison_df.to_csv(DATA_DIR / 'comparison.csv', index=False)
    
    return f"Generated comparison CSV with {len(rows)} rows. Saved to {DATA_DIR / 'comparison.csv'}"

@tool
def export_narrative(run_key: str) -> str:
    """
    Send themes and taxonomy to Mistral to draft a 500-word Section 7 literature review.
    References B&C phases.
    """
    run_dir = DATA_DIR / run_key
    
    with open(run_dir / 'themes.json', 'r') as f:
        themes = json.load(f)
        
    with open(run_dir / 'taxonomy_map.json', 'r') as f:
        taxonomy = json.load(f)
        
    llm = ChatMistralAI(model='mistral-small-latest', temperature=0.7)
    
    prompt = PromptTemplate(
        template="""You are writing Section 7 (Literature Review) for an academic paper on IS research trends.
Themes discovered:
{themes}
Taxonomy mapping:
{taxonomy}
Write a 500-word narrative that:
1. References the Bandara & Cresswell (B&C) framework phases
2. Highlights dominant research themes
3. Notes novel or emerging areas
4. Connects to PAJAIS categories
5. Uses academic tone
Output only the prose, no preamble.""",
        input_variables=['themes', 'taxonomy']
    )
    
    themes_text = json.dumps([{
        'theme_name': t['theme_name'],
        'paper_count': t['paper_count']
    } for t in themes], indent=2)
    
    taxonomy_text = json.dumps([{
        'theme_name': m['theme_name'],
        'pajais_match': m['pajais_match'],
        'is_novel': m['is_novel']
    } for m in taxonomy], indent=2)
    
    chain = prompt | llm
    narrative = chain.invoke({'themes': themes_text, 'taxonomy': taxonomy_text})
    narrative_text = narrative.content
    
    with open(run_dir / 'narrative.txt', 'w') as f:
        f.write(narrative_text)
        
    return f"Generated 500-word narrative. Saved to {run_dir / 'narrative.txt'}"