"""
Intel Reddit AI Analyzer - Hugging Face Space
DistilBertForSequenceClassification (6 layers, 768 dim)
Labels: 0=negative, 1=neutral, 2=positive
Endpoints match client.ts: /analyze_reddit_content, /analyze_sentiment, /deep_analyze, /predict
"""

import asyncio
import json
import logging
import os
import signal
import sys
import warnings

import gradio as gr
import numpy as np
import spacy
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# ── Logging ──────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Completely suppress asyncio cleanup warnings that are harmless in containers
# These warnings occur during normal container shutdown and don't affect functionality
warnings.filterwarnings("ignore", message=".*Invalid file descriptor.*")
warnings.filterwarnings("ignore", message=".*Exception ignored in.*")
warnings.filterwarnings("ignore", category=UserWarning, module="asyncio")

# Custom exception hook to suppress asyncio cleanup errors
def custom_excepthook(exc_type, exc_value, exc_traceback):
    if "Invalid file descriptor" in str(exc_value) and "asyncio" in str(exc_traceback):
        return  # Suppress these specific asyncio cleanup errors
    sys.__excepthook__(exc_type, exc_value, exc_traceback)

sys.excepthook = custom_excepthook

# ── Labels ────────────────────────────────────────────────────────────────────
# classes.npy has typos ('netural sentiment') so we force clean labels.
# config.json id2label is generic (LABEL_0/1/2), order confirmed from training:
LABELS = ["negative", "neutral", "positive"]  # index 0, 1, 2

# ── Load model ────────────────────────────────────────────────────────────────
MODEL_PATH = "/app/models" if os.path.exists("/app/models") else "./models"
logger.info(f"Loading model from {MODEL_PATH} ...")

from transformers import BertTokenizerFast, DistilBertConfig, DistilBertForSequenceClassification
from transformers import PreTrainedTokenizerFast

_tok_json = os.path.join(MODEL_PATH, "tokenizer.json")
_vocab_txt = os.path.join(MODEL_PATH, "vocab.txt")
if os.path.exists(_tok_json):
    tokenizer = PreTrainedTokenizerFast(tokenizer_file=_tok_json)
elif os.path.exists(_vocab_txt):
    tokenizer = BertTokenizerFast(vocab_file=_vocab_txt)
else:
    raise FileNotFoundError(f"No tokenizer files found in {MODEL_PATH}")

_config = DistilBertConfig.from_json_file(os.path.join(MODEL_PATH, "config.json"))
model = DistilBertForSequenceClassification(_config)

# Fix padding token issue - ensure it's properly set
if tokenizer.pad_token is None:
    if tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token
    else:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})

# Resize model embeddings to accommodate new padding token if needed
if tokenizer.pad_token not in tokenizer.get_vocab():
    model.resize_token_embeddings(len(tokenizer))

_safetensors = os.path.join(MODEL_PATH, "model.safetensors")
_pytorch_bin = os.path.join(MODEL_PATH, "pytorch_model.bin")
if os.path.exists(_safetensors):
    from safetensors.torch import load_file as _load_safetensors
    _state_dict = _load_safetensors(_safetensors)
elif os.path.exists(_pytorch_bin):
    _state_dict = torch.load(_pytorch_bin, map_location="cpu", weights_only=True)
else:
    raise FileNotFoundError(f"No model weights found in {MODEL_PATH}")

_key_map = {"gamma": "weight", "beta": "bias"}
_state_dict = {
    ".".join(_key_map.get(part, part) for part in k.split(".")): v
    for k, v in _state_dict.items()
}

model.load_state_dict(_state_dict)
model.eval()
logger.info("Model loaded.")

# ── Load spaCy for location extraction ───────────────────────────────────────
try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    from spacy.cli import download
    download("en_core_web_sm")
    nlp = spacy.load("en_core_web_sm")
logger.info("spaCy loaded.")


# ── Core helpers ──────────────────────────────────────────────────────────────

def _predict_probs(text: str):
    """Run DistilBERT forward pass, return (label_str, confidence, all_probs_dict)."""
    inputs = tokenizer(
        text, padding="max_length", truncation=True, max_length=512, return_tensors="pt"
    )
    inputs.pop("token_type_ids", None)  # DistilBERT has no token_type_ids

    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.nn.functional.softmax(logits, dim=-1)[0]

    idx = int(torch.argmax(probs).item())
    label = LABELS[idx] if idx < len(LABELS) else f"LABEL_{idx}"
    confidence = float(probs[idx].item())
    all_probs = {LABELS[i]: float(probs[i].item()) for i in range(len(LABELS))}
    return label, confidence, all_probs, idx


def get_gradient_word_importance(text: str, prediction_idx: int):
    """
    Gradient × embedding saliency with proper sub-word merging and noise filtering.
    Single backward pass — fast, model-derived, no heuristics.
    Returns list of {word, importance, sentiment_contribution} sorted by |importance|.
    """
    import re
    try:
        inputs = tokenizer(
            text, padding=True, truncation=True, max_length=512, return_tensors="pt"
        )
        inputs.pop("token_type_ids", None)
        tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

        # DistilBERT embedding layer
        embeddings = model.distilbert.embeddings(inputs["input_ids"])
        embeddings = embeddings.detach().requires_grad_(True)

        outputs = model(
            inputs_embeds=embeddings,
            attention_mask=inputs.get("attention_mask"),
        )
        model.zero_grad()
        outputs.logits[0][prediction_idx].backward()

        # importance per token = sum_d(grad_d × emb_d)
        importance = (embeddings.grad * embeddings).sum(dim=-1)[0]  # (seq_len,)

        SPECIAL = {"[CLS]", "[SEP]", "[PAD]", "[MASK]", "[UNK]"}
        STOPWORDS = {
            'the', 'a', 'an', 'is', 'it', 'in', 'on', 'at', 'to', 'for', 'of', 'and',
            'or', 'but', 'this', 'that', 'with', 'are', 'was', 'be', 'as', 'by',
            'from', 'its', 'https', 'http', 'www', 'com', 'i', 'we', 'my', 'he',
            'she', 'they', 'you', 'not', 'so', 'if', 'do', 'no', 'up', 'out', 'what',
            'how', 'just', 'get', 'has', 'had', 'have', 'about', 'been', 'which',
            'their', 'there', 'can', 'will', 'would', 'one', 'more', 'like', 'all',
        }

        # --- Merge sub-word tokens (## prefix = continuation) ---
        merged = []  # list of [word_str, importance_sum]
        for token, imp in zip(tokens, importance.tolist()):
            if token in SPECIAL:
                continue
            is_continuation = token.startswith("##")
            clean = token.replace("##", "").strip()
            if not clean:
                continue
            if is_continuation and merged:
                merged[-1][0] += clean
                merged[-1][1] += imp
            else:
                merged.append([clean, imp])

        # --- Filter noise ---
        results = []
        for word, imp in merged:
            if word.lower() in STOPWORDS:
                continue
            if not any(c.isalpha() for c in word):  # must contain a letter
                continue
            if len(word) < 2:
                continue
            if re.search(r'https?|www\.', word, re.I):  # skip URL fragments
                continue
            results.append({
                "word": word,
                "importance": round(imp, 4),
                "sentiment_contribution": "positive" if imp > 0 else "negative",
            })

        results.sort(key=lambda x: abs(x["importance"]), reverse=True)
        return results[:5]

    except Exception as e:
        logger.error(f"Gradient importance error: {e}")
        return []


def extract_locations(texts: list[str]) -> list[str]:
    """Extract GPE/LOC entities from a list of texts using spaCy."""
    try:
        combined = " ".join(texts[:25])
        doc = nlp(combined)
        locs = {ent.text for ent in doc.ents if ent.label_ in ("GPE", "LOC")}
        return list(locs) if locs else ["No specific locations detected"]
    except Exception as e:
        logger.error(f"Location extraction error: {e}")
        return ["Location detection failed"]


# ── Gradio endpoint functions ─────────────────────────────────────────────────

def analyze_reddit_content(posts_json: str, comments_json: str) -> dict:
    """
    Batch analysis of posts and comments.
    Called by client.ts → analyzeWithHuggingFace() via /analyze_reddit_content
    Input:  JSON strings of post/comment arrays
    Output: dict with post_sentiments, comment_sentiments, breakdown, locations, etc.
    """
    try:
        posts = json.loads(posts_json) if posts_json else []
        comments = json.loads(comments_json) if comments_json else []
    except json.JSONDecodeError as e:
        return {"error": f"Invalid JSON: {e}"}

    all_texts = []
    post_sentiments = []
    comment_sentiments = []

    for p in posts:
        title = p.get("title", "")
        body = p.get("selftext", "")
        text = f"{title} {body}".strip()
        if not text:
            # Skip empty posts - no text to analyze
            post_sentiments.append({
                "sentiment": "neutral",
                "confidence": 0.0,
                "all_probabilities": {l: 0.0 for l in LABELS},
                "text_preview": "",
                "word_importance": [],
            })
            continue
        all_texts.append(text)

        label, conf, all_probs, pred_idx = _predict_probs(text)
        word_importance = get_gradient_word_importance(text, pred_idx)
        post_sentiments.append(
            {
                "sentiment": label,
                "confidence": round(conf, 4),
                "all_probabilities": {k: round(v, 4) for k, v in all_probs.items()},
                "text_preview": text[:150],
                "word_importance": word_importance,
            }
        )

    for c in comments:
        text = c.get("body", "").strip()
        if not text:
            # Skip empty comments - no text to analyze
            comment_sentiments.append({
                "sentiment": "neutral",
                "confidence": 0.0,
                "all_probabilities": {l: 0.0 for l in LABELS},
                "text_preview": "",
                "word_importance": [],
            })
            continue
        all_texts.append(text)

        label, conf, all_probs, pred_idx = _predict_probs(text)
        word_importance = get_gradient_word_importance(text, pred_idx)
        comment_sentiments.append(
            {
                "sentiment": label,
                "confidence": round(conf, 4),
                "all_probabilities": {k: round(v, 4) for k, v in all_probs.items()},
                "text_preview": text[:150],
                "word_importance": word_importance,
            }
        )

    # Aggregate sentiment breakdown across posts + comments
    all_sentiments = post_sentiments + comment_sentiments
    counts = {lbl: 0 for lbl in LABELS}
    for item in all_sentiments:
        s = item["sentiment"]
        if s in counts:
            counts[s] += 1

    total = len(all_sentiments) or 1
    breakdown = {k: round(v / total, 4) for k, v in counts.items()}
    dominant = max(counts, key=counts.get) if all_sentiments else "neutral"

    return {
        "post_sentiments": post_sentiments,
        "comment_sentiments": comment_sentiments,
        "post_count": len(post_sentiments),
        "comment_count": len(comment_sentiments),
        "sentiment_breakdown": breakdown,
        "locations": extract_locations(all_texts),
        "dominant_sentiment": dominant,
    }


def analyze_sentiment(text: str) -> dict:
    """
    Single text sentiment analysis.
    Called by client.ts → analyzeSingleText() via /analyze_sentiment
    """
    if not text or not text.strip():
        return {
            "sentiment": "neutral",
            "confidence": 0.0,
            "all_probabilities": {l: 0.0 for l in LABELS},
            "text_preview": "",
        }

    label, conf, all_probs, _ = _predict_probs(text)
    return {
        "sentiment": label,
        "confidence": round(conf, 4),
        "all_probabilities": {k: round(v, 4) for k, v in all_probs.items()},
        "text_preview": text[:150],
    }


def deep_analyze(text: str) -> dict:
    """
    Deep analysis with gradient-based word importance.
    Called by client.ts → analyzeDeep() via /deep_analyze
    """
    if not text or not text.strip():
        return {
            "text": "",
            "overall_sentiment": "neutral",
            "confidence": 0.0,
            "word_importance": [],
            "explanation": "No text provided.",
        }

    label, conf, all_probs, pred_idx = _predict_probs(text)
    word_importance = get_gradient_word_importance(text, pred_idx)

    pos_words = [w["word"] for w in word_importance if w["sentiment_contribution"] == "positive"]
    neg_words = [w["word"] for w in word_importance if w["sentiment_contribution"] == "negative"]
    explanation = f"Classified as {label} with {conf:.0%} confidence."
    if pos_words:
        explanation += f" Supporting words: {', '.join(pos_words[:3])}."
    if neg_words:
        explanation += f" Opposing words: {', '.join(neg_words[:3])}."

    return {
        "text": text[:200],
        "overall_sentiment": label,
        "confidence": round(conf, 4),
        "word_importance": word_importance,
        "explanation": explanation,
    }


def predict(posts_json: str, comments_json: str, deep_text: str = "") -> dict:
    """
    Unified endpoint.
    Called by client.ts → predict() via /predict
    If deep_text is provided → deep analysis, otherwise → batch analysis.
    """
    if deep_text and deep_text.strip():
        return deep_analyze(deep_text)
    return analyze_reddit_content(posts_json, comments_json)


# ── Gradio UI ─────────────────────────────────────────────────────────────────

with gr.Blocks(title="Intel Reddit AI Analyzer") as demo:
    gr.Markdown(
        """
# 🔍 Intel Reddit AI Analyzer
**AI-powered sentiment analysis and location extraction for Reddit OSINT investigations**

Model: DistilBERT fine-tuned for sentiment classification (Negative / Neutral / Positive)

> **Note**: Runs on CPU (free tier). First request may take 10–30 seconds to load.
        """
    )

    with gr.Tab("Batch Analysis"):
        gr.Markdown("Analyze multiple Reddit posts and comments at once.")
        posts_input = gr.Textbox(
            label="Posts JSON",
            placeholder='[{"title": "...", "selftext": "..."}]',
            lines=4,
        )
        comments_input = gr.Textbox(
            label="Comments JSON",
            placeholder='[{"body": "..."}]',
            lines=4,
        )
        batch_btn = gr.Button("Analyze Batch")
        batch_output = gr.JSON(label="Results")
        batch_btn.click(
            fn=analyze_reddit_content,
            inputs=[posts_input, comments_input],
            outputs=batch_output,
            api_name="analyze_reddit_content",
        )

    with gr.Tab("Single Text"):
        gr.Markdown("Analyze a single piece of text.")
        single_input = gr.Textbox(label="Text", lines=3)
        single_btn = gr.Button("Analyze")
        single_output = gr.JSON(label="Results")
        single_btn.click(
            fn=analyze_sentiment,
            inputs=single_input,
            outputs=single_output,
            api_name="analyze_sentiment",
        )

    with gr.Tab("Deep Analysis"):
        gr.Markdown("Word-level sentiment analysis using gradient saliency.")
        deep_input = gr.Textbox(label="Text", lines=3)
        deep_btn = gr.Button("Deep Analyze")
        deep_output = gr.JSON(label="Results")
        deep_btn.click(
            fn=deep_analyze,
            inputs=deep_input,
            outputs=deep_output,
            api_name="deep_analyze",
        )

    with gr.Tab("Unified Predict"):
        gr.Markdown("Batch or deep analysis in one endpoint.")
        p_posts = gr.Textbox(label="Posts JSON", lines=3)
        p_comments = gr.Textbox(label="Comments JSON", lines=3)
        p_deep = gr.Textbox(label="Deep text (optional — leave blank for batch)", lines=2)
        p_btn = gr.Button("Predict")
        p_output = gr.JSON(label="Results")
        p_btn.click(
            fn=predict,
            inputs=[p_posts, p_comments, p_deep],
            outputs=p_output,
            api_name="predict",
        )

# ── Launch with asyncio error suppression ───────────────────────────────────────
if __name__ == "__main__":
    try:
        # Suppress asyncio warnings that are harmless in container environments
        import asyncio
        import warnings
        
        # These warnings are harmless and occur during normal container shutdown
        warnings.filterwarnings("ignore", category=UserWarning, module="asyncio")
        
        # Set event loop policy for better container compatibility
        if sys.platform.startswith('linux'):
            asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy())
        
        # Launch with SSR disabled to prevent asyncio event loop issues
        # SSR mode can cause event loop cleanup problems in containers
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            show_error=True,
            quiet=False,
            ssr_mode=False
        )
    except KeyboardInterrupt:
        logger.info("Application interrupted")
    except Exception as e:
        logger.error(f"Fatal error: {e}")
        sys.exit(1)