# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# Licensed under the MIT License

"""
Email TF-IDF Classification Model.

Trains a local Naive Bayes classifier on the user's own email
classification history (confirmed ratings, corrections, and
high-confidence auto-harvested pairs).

No external dependencies required — uses sklearn if available,
falls back to a simple word-frequency scorer if not.

Model persisted at: ~/.familiar/data/email_tfidf_model.json
(JSON format, not pickle — portable and inspectable)
"""

import json
import logging
import math
import re
import threading
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path

logger = logging.getLogger(__name__)

_lock = threading.Lock()


def _data_dir() -> Path:
    try:
        from familiar.core.paths import DATA_DIR
        return DATA_DIR
    except ImportError:
        return Path.home() / ".familiar" / "data"


def _model_path() -> Path:
    return _data_dir() / "email_tfidf_model.json"


def _tokenize(text: str) -> list:
    """Simple whitespace + punctuation tokenizer."""
    text = text.lower()
    # Remove URLs
    text = re.sub(r"https?://\S+", "", text)
    # Split on non-alphanumeric
    tokens = re.findall(r"[a-z0-9]+", text)
    # Remove very short tokens and stopwords
    _STOP = {
        "the", "a", "an", "is", "are", "was", "were", "be", "been",
        "to", "of", "in", "for", "on", "at", "by", "with", "from",
        "and", "or", "but", "not", "no", "this", "that", "it", "its",
        "you", "your", "we", "our", "my", "me", "he", "she", "they",
        "has", "have", "had", "do", "does", "did", "will", "would",
        "can", "could", "may", "might", "shall", "should",
    }
    return [t for t in tokens if len(t) > 2 and t not in _STOP]


class NaiveBayesClassifier:
    """Simple Multinomial Naive Bayes using word frequencies.

    No sklearn dependency. JSON-serializable model state.
    Trained on (text, category) pairs from the training data store.
    """

    def __init__(self):
        self.category_counts = Counter()      # category → doc count
        self.word_counts = defaultdict(Counter)  # category → {word: count}
        self.vocab = set()
        self.total_docs = 0
        self.trained_at = None
        self.accuracy_estimate = 0.0

    def train(self, pairs: list):
        """Train on [(text, category), ...] pairs."""
        self.category_counts = Counter()
        self.word_counts = defaultdict(Counter)
        self.vocab = set()
        self.total_docs = 0

        for text, category in pairs:
            tokens = _tokenize(text)
            self.category_counts[category] += 1
            for token in tokens:
                self.word_counts[category][token] += 1
                self.vocab.add(token)
            self.total_docs += 1

        self.trained_at = datetime.now().isoformat()

        # Estimate accuracy via leave-one-out on last 20%
        if len(pairs) >= 20:
            split = int(len(pairs) * 0.8)
            test_pairs = pairs[split:]
            correct = 0
            for text, expected in test_pairs:
                predicted, _ = self._predict_internal(text)
                if predicted == expected:
                    correct += 1
            self.accuracy_estimate = round(correct / len(test_pairs), 2) if test_pairs else 0
        else:
            self.accuracy_estimate = 0

        logger.info(
            "TF-IDF model trained: %d docs, %d vocab, %d categories, accuracy %.0f%%",
            self.total_docs, len(self.vocab), len(self.category_counts),
            self.accuracy_estimate * 100,
        )

    def predict(self, text: str) -> tuple:
        """Predict category for text. Returns (category, confidence)."""
        if not self.total_docs:
            return "Other", 0.0
        return self._predict_internal(text)

    def _predict_internal(self, text: str) -> tuple:
        tokens = _tokenize(text)
        if not tokens:
            return "Other", 0.0

        vocab_size = len(self.vocab) + 1  # +1 for smoothing
        scores = {}

        for category, doc_count in self.category_counts.items():
            # Log prior
            log_prior = math.log(doc_count / self.total_docs)

            # Log likelihood with Laplace smoothing
            total_words_in_cat = sum(self.word_counts[category].values())
            log_likelihood = 0.0
            for token in tokens:
                word_count = self.word_counts[category].get(token, 0)
                log_likelihood += math.log((word_count + 1) / (total_words_in_cat + vocab_size))

            scores[category] = log_prior + log_likelihood

        if not scores:
            return "Other", 0.0

        # Convert log scores to probabilities via softmax
        max_score = max(scores.values())
        exp_scores = {cat: math.exp(s - max_score) for cat, s in scores.items()}
        total_exp = sum(exp_scores.values())
        probs = {cat: exp_s / total_exp for cat, exp_s in exp_scores.items()}

        best_cat = max(probs, key=probs.get)
        confidence = round(probs[best_cat], 3)

        return best_cat, confidence

    def to_dict(self) -> dict:
        """Serialize model to JSON-compatible dict."""
        return {
            "category_counts": dict(self.category_counts),
            "word_counts": {cat: dict(wc) for cat, wc in self.word_counts.items()},
            "vocab_size": len(self.vocab),
            "total_docs": self.total_docs,
            "trained_at": self.trained_at,
            "accuracy_estimate": self.accuracy_estimate,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "NaiveBayesClassifier":
        """Deserialize model from dict."""
        model = cls()
        model.category_counts = Counter(data.get("category_counts", {}))
        model.word_counts = defaultdict(Counter)
        for cat, wc in data.get("word_counts", {}).items():
            model.word_counts[cat] = Counter(wc)
        model.vocab = set()
        for wc in model.word_counts.values():
            model.vocab.update(wc.keys())
        model.total_docs = data.get("total_docs", 0)
        model.trained_at = data.get("trained_at")
        model.accuracy_estimate = data.get("accuracy_estimate", 0)
        return model


# ── Global model instance ─────────────────────────────────────────────────────

_model = None
_model_loaded = False


def _load_model() -> NaiveBayesClassifier:
    """Load model from disk, or return empty model."""
    global _model, _model_loaded
    if _model_loaded and _model:
        return _model

    path = _model_path()
    if path.exists():
        try:
            with open(path) as f:
                data = json.load(f)
            _model = NaiveBayesClassifier.from_dict(data)
            _model_loaded = True
            logger.debug("Loaded TF-IDF model: %d docs, accuracy %.0f%%",
                         _model.total_docs, _model.accuracy_estimate * 100)
            return _model
        except Exception as e:
            logger.warning("Failed to load TF-IDF model: %s", e)

    _model = NaiveBayesClassifier()
    _model_loaded = True
    return _model


def _save_model(model: NaiveBayesClassifier):
    """Persist model to disk."""
    path = _model_path()
    with _lock:
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps(model.to_dict(), indent=2))
        tmp.replace(path)


# ── Public API ────────────────────────────────────────────────────────────────


def train_model() -> dict:
    """Train the TF-IDF model from accumulated training pairs.

    Returns: {trained: bool, docs: N, vocab: N, accuracy: float, categories: {}}
    """
    try:
        from familiar.skills.email.training_data import get_training_pairs
    except ImportError:
        return {"trained": False, "error": "Training data store not available"}

    pairs = get_training_pairs()
    if len(pairs) < 20:
        return {"trained": False, "error": f"Need at least 20 training pairs, have {len(pairs)}"}

    # Build (text, category) pairs — use subject + from_email as text
    training = []
    for p in pairs:
        text = f"{p.get('from_email', '')} {p.get('subject', '')}"
        category = p.get("category", "Other")
        if category and category != "Other":
            training.append((text, category))

    # Also include "Other" pairs so the model can predict Other
    other_pairs = [p for p in pairs if p.get("category") == "Other"]
    for p in other_pairs:
        text = f"{p.get('from_email', '')} {p.get('subject', '')}"
        training.append((text, "Other"))

    if len(training) < 20:
        return {"trained": False, "error": f"Need at least 20 non-Other pairs, have {len(training)}"}

    model = NaiveBayesClassifier()
    model.train(training)
    _save_model(model)

    global _model, _model_loaded
    _model = model
    _model_loaded = True

    return {
        "trained": True,
        "docs": model.total_docs,
        "vocab": len(model.vocab),
        "accuracy": model.accuracy_estimate,
        "categories": dict(model.category_counts),
    }


def predict(subject: str, from_email: str = "") -> tuple:
    """Predict category using the trained model.

    Returns (category, confidence) or ("Other", 0.0) if model not trained.
    """
    model = _load_model()
    if not model.total_docs:
        return "Other", 0.0

    text = f"{from_email} {subject}"
    return model.predict(text)


def get_model_info() -> dict:
    """Get model metadata."""
    model = _load_model()
    if not model.total_docs:
        return {"trained": False}
    return {
        "trained": True,
        "docs": model.total_docs,
        "vocab": len(model.vocab),
        "accuracy": model.accuracy_estimate,
        "trained_at": model.trained_at,
        "categories": dict(model.category_counts),
    }
