"""
task2_module.py
───────────────
BERT Semantic Encoder using all-mpnet-base-v2.
Depends on: task1_module.py

Usage:
    from task2_module import load_bert_model, bert_transform, bert_batch_transform
"""
import logging


import os
import numpy as np
from sentence_transformers import SentenceTransformer

# Import flatten from Task 1 (must be in same directory)
from task1_module import flatten


# ────────────────────────────────────────────────────
# SECTION A: Model Loading
# ────────────────────────────────────────────────────

def load_bert_model(model_path: str = 'bert_model_saved', device: str = None):
    """
    Load BERT model from a saved folder OR download from HuggingFace.

    Parameters
    ----------
    model_path : str
        Path to saved folder (from Task 2 output) OR
        HuggingFace model name ('sentence-transformers/all-mpnet-base-v2')
    device : str, optional
        'cuda' or 'cpu'. Auto-detects if None.

    Returns
    -------
    SentenceTransformer
    """
    import torch
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    logging.getLogger('task2').info('loading BERT model', extra={'f': {'path': model_path, 'device': device}})
    model = SentenceTransformer(model_path, device=device)
    logging.getLogger('task2').info('BERT loaded', extra={'f': {'dims': model.get_sentence_embedding_dimension()}})
    return model


# ────────────────────────────────────────────────────
# SECTION B: Encoding Functions
# ────────────────────────────────────────────────────

def _truncate(text: str, max_words: int = 400) -> str:
    """Keep first 200 + last 200 words if text is too long."""
    words = text.split()
    if len(words) > max_words:
        return ' '.join(words[:200] + words[-200:])
    return text


def bert_encode(text: str, model: SentenceTransformer) -> np.ndarray:
    """
    Encode a pre-cleaned string → 768-dim L2-normalised vector.
    """
    if not isinstance(text, str) or len(text.strip()) == 0:
        return np.zeros(model.get_sentence_embedding_dimension())
    return model.encode(
        _truncate(text),
        normalize_embeddings=True,
        show_progress_bar=False,
        convert_to_numpy=True
    )


def bert_transform(input_data, model: SentenceTransformer, mode: str = 'auto') -> np.ndarray:
    """
    Full pipeline: raw input → flatten (Task1) → BERT encode → 768-dim vector.

    Parameters
    ----------
    input_data : dict | pd.Series | str
    model      : loaded SentenceTransformer
    mode       : 'json' | 'csv' | 'text' | 'auto'

    Returns
    -------
    np.ndarray : shape (768,)
    """
    flat = flatten(input_data, mode=mode)
    return bert_encode(flat, model=model)


def bert_batch_transform(text_list: list, model: SentenceTransformer,
                          batch_size: int = 32) -> np.ndarray:
    """
    Encode a list of already-flattened strings in batches.

    Returns
    -------
    np.ndarray : shape (N, 768)
    """
    truncated = [_truncate(t) for t in text_list]
    return model.encode(
        truncated,
        batch_size=batch_size,
        normalize_embeddings=True,
        show_progress_bar=True,
        convert_to_numpy=True
    )