"""CRF Sequence Labeler for Japanese email line classification.

Uses python-crfsuite to predict labels for each line of an email:
- GREETING: Opening formulas
- BODY: Substantive content
- CLOSING: Closing formulas
- SIGNATURE: Signature block lines
- QUOTE: Quoted content
- OTHER: Visual dividers, headers, noise, unclassifiable
"""

import logging
from dataclasses import dataclass
from importlib import resources
from pathlib import Path
from typing import Literal

import pycrfsuite

from yomail.pipeline.features import ExtractedFeatures, LineFeatures

logger = logging.getLogger(__name__)


def get_default_model_path() -> Path:
    """Get path to the bundled default model.

    Returns:
        Path to the bundled email_body.crfsuite model.

    Raises:
        FileNotFoundError: If the bundled model is not found.
    """
    data_files = resources.files("yomail.data")
    model_file = data_files.joinpath("email_body.crfsuite")
    # Use as_file to get a concrete path (needed for pycrfsuite)
    with resources.as_file(model_file) as path:
        if not path.exists():
            raise FileNotFoundError("Bundled model not found: email_body.crfsuite")
        # Return a copy of the path since we're outside the context manager
        return Path(path)


# Label definitions
Label = Literal["GREETING", "BODY", "CLOSING", "SIGNATURE", "QUOTE", "OTHER"]

LABELS: tuple[Label, ...] = ("GREETING", "BODY", "CLOSING", "SIGNATURE", "QUOTE", "OTHER")


@dataclass(frozen=True, slots=True)
class LabeledLine:
    """A line with its predicted label and confidence.

    Attributes:
        text: Original line text.
        label: Predicted label (one of LABELS).
        confidence: Marginal probability for the predicted label (0.0 to 1.0).
        label_probabilities: Marginal probabilities for all labels.
    """

    text: str
    label: Label
    confidence: float
    label_probabilities: dict[Label, float]


@dataclass(frozen=True, slots=True)
class SequenceLabelingResult:
    """Result of CRF sequence labeling.

    Attributes:
        labeled_lines: Tuple of LabeledLine objects, one per input line.
        sequence_probability: Probability of the entire predicted label sequence.
    """

    labeled_lines: tuple[LabeledLine, ...]
    sequence_probability: float


def _features_to_dict(features: LineFeatures, idx: int, total_lines: int) -> dict[str, str | float | bool]:
    """Convert LineFeatures to a feature dictionary for CRFsuite.

    CRFsuite accepts features as dict[str, value] where value can be:
    - str for categorical features
    - float for numeric features
    - bool for boolean features

    Args:
        features: LineFeatures for a single line.
        idx: Line index (for BOS/EOS markers).
        total_lines: Total number of lines.

    Returns:
        Feature dictionary for CRFsuite.
    """
    feat: dict[str, str | float | bool] = {}

    # Special markers for beginning/end of sequence
    if idx == 0:
        feat["BOS"] = True
    if idx == total_lines - 1:
        feat["EOS"] = True

    # Positional features (numeric)
    feat["pos_norm"] = features.position_normalized
    feat["pos_rev"] = features.position_reverse
    feat["lines_from_start"] = float(features.lines_from_start)
    feat["lines_from_end"] = float(features.lines_from_end)
    feat["pos_rel_first_quote"] = features.position_rel_first_quote
    feat["pos_rel_last_quote"] = features.position_rel_last_quote

    # Content features (numeric)
    feat["line_length"] = float(features.line_length)
    feat["kanji_ratio"] = features.kanji_ratio
    feat["hiragana_ratio"] = features.hiragana_ratio
    feat["katakana_ratio"] = features.katakana_ratio
    feat["ascii_ratio"] = features.ascii_ratio
    feat["digit_ratio"] = features.digit_ratio
    feat["symbol_ratio"] = features.symbol_ratio
    feat["leading_ws"] = float(features.leading_whitespace)
    feat["trailing_ws"] = float(features.trailing_whitespace)

    # Whitespace context features
    feat["blank_lines_before"] = float(features.blank_lines_before)
    feat["blank_lines_after"] = float(features.blank_lines_after)

    # Structural features
    feat["quote_depth"] = float(features.quote_depth)
    feat["is_forward_reply_header"] = features.is_forward_reply_header
    feat["preceded_by_delimiter"] = features.preceded_by_delimiter
    feat["is_delimiter"] = features.is_delimiter

    # Pattern flags (boolean)
    feat["is_greeting"] = features.is_greeting
    feat["is_closing"] = features.is_closing
    feat["has_contact_info"] = features.has_contact_info
    feat["has_company_pattern"] = features.has_company_pattern
    feat["has_position_pattern"] = features.has_position_pattern
    feat["has_name_pattern"] = features.has_name_pattern
    feat["is_visual_separator"] = features.is_visual_separator
    feat["has_meta_discussion"] = features.has_meta_discussion
    feat["is_inside_quotation_marks"] = features.is_inside_quotation_marks

    # Contextual features (numeric)
    feat["ctx_greeting_count"] = float(features.context_greeting_count)
    feat["ctx_closing_count"] = float(features.context_closing_count)
    feat["ctx_contact_count"] = float(features.context_contact_count)
    feat["ctx_quote_count"] = float(features.context_quote_count)
    feat["ctx_separator_count"] = float(features.context_separator_count)

    # Derived categorical features for stronger signal
    if features.quote_depth > 0:
        feat["quote_depth_cat"] = "quoted"
    else:
        feat["quote_depth_cat"] = "unquoted"

    # Position buckets for categorical features
    if features.position_normalized < 0.1:
        feat["pos_bucket"] = "start"
    elif features.position_normalized < 0.3:
        feat["pos_bucket"] = "early"
    elif features.position_normalized < 0.7:
        feat["pos_bucket"] = "middle"
    elif features.position_normalized < 0.9:
        feat["pos_bucket"] = "late"
    else:
        feat["pos_bucket"] = "end"

    # Character composition bucket (content lines only, no blank option)
    if features.ascii_ratio > 0.8:
        feat["char_type"] = "ascii_heavy"
    elif features.kanji_ratio + features.hiragana_ratio > 0.7:
        feat["char_type"] = "japanese_heavy"
    else:
        feat["char_type"] = "mixed"

    # Bracket features
    feat["in_bracketed_section"] = features.in_bracketed_section
    feat["bracket_has_signature_patterns"] = features.bracket_has_signature_patterns

    # Derived categorical for brackets
    if features.in_bracketed_section:
        feat["bracket_cat"] = "bracketed"
    else:
        feat["bracket_cat"] = "unbracketed"

    return feat


def _extract_feature_sequence(
    extracted: ExtractedFeatures,
    texts: tuple[str, ...],
) -> list[dict[str, str | float | bool]]:
    """Extract feature sequence for all lines.

    Args:
        extracted: ExtractedFeatures from the FeatureExtractor.
        texts: Original text lines (for reference, not used in features).

    Returns:
        List of feature dictionaries, one per line.
    """
    total_lines = extracted.total_lines
    return [
        _features_to_dict(line_features, idx, total_lines)
        for idx, line_features in enumerate(extracted.line_features)
    ]


class CRFSequenceLabeler:
    """CRF-based sequence labeler for email line classification.

    Uses python-crfsuite for fast, lightweight CRF inference.
    Supports loading pre-trained models and making predictions
    with per-label marginal probabilities.
    """

    def __init__(self, model_path: Path | str | None = None, use_default: bool = True) -> None:
        """Initialize the CRF labeler.

        Args:
            model_path: Path to a trained CRF model file.
                If None and use_default is True, loads the bundled model.
            use_default: If True and model_path is None, load the bundled model.
                Set to False to create an unloaded labeler.
        """
        self._tagger: pycrfsuite.Tagger | None = None
        self._model_path: Path | None = None
        self._resource_context = None  # Keep resource context alive

        if model_path is not None:
            self.load_model(model_path)
        elif use_default:
            self._load_default_model()

    def load_model(self, model_path: Path | str) -> None:
        """Load a trained CRF model.

        Args:
            model_path: Path to the model file.

        Raises:
            FileNotFoundError: If the model file does not exist.
            RuntimeError: If the model cannot be loaded.
        """
        path = Path(model_path)
        if not path.exists():
            raise FileNotFoundError(f"Model file not found: {path}")

        self._tagger = pycrfsuite.Tagger()
        try:
            self._tagger.open(str(path))
        except Exception as exc:
            self._tagger = None
            raise RuntimeError(f"Failed to load CRF model: {exc}") from exc

        self._model_path = path
        logger.info("Loaded CRF model from %s", path)

    def _load_default_model(self) -> None:
        """Load the bundled default model."""
        data_files = resources.files("yomail.data")
        model_file = data_files.joinpath("email_body.crfsuite")
        # Keep context manager alive for duration of object
        self._resource_context = resources.as_file(model_file)
        path = self._resource_context.__enter__()

        self._tagger = pycrfsuite.Tagger()
        try:
            self._tagger.open(str(path))
        except Exception as exc:
            self._tagger = None
            self._resource_context.__exit__(None, None, None)
            self._resource_context = None
            raise RuntimeError(f"Failed to load bundled CRF model: {exc}") from exc

        self._model_path = path
        logger.info("Loaded bundled CRF model")

    @property
    def is_loaded(self) -> bool:
        """Whether a model is currently loaded."""
        return self._tagger is not None

    @property
    def labels(self) -> tuple[str, ...]:
        """Get the labels known by the loaded model."""
        if self._tagger is None:
            return LABELS
        return tuple(self._tagger.labels())

    def _fix_impossible_transitions(
        self, labels: list[str], extracted: ExtractedFeatures
    ) -> list[str]:
        """Fix impossible label transitions that the CRF may produce.

        The CRF can predict impossible transitions like SIGNATURE → CLOSING
        because it never saw them in training (weight = 0, not negative).
        This post-processing step enforces structural constraints.

        Rules:
        - SIGNATURE → CLOSING is impossible (relabel CLOSING as SIGNATURE)
        - Separator lines (is_delimiter or is_visual_separator) cannot be CLOSING
        """
        if len(labels) < 2:
            return labels

        fixed = labels.copy()
        in_signature = False

        for i in range(len(fixed)):
            line_features = extracted.line_features[i]

            # Rule: Separator lines cannot be CLOSING
            # They should inherit from context (SIGNATURE if after body, BODY if in info block)
            if fixed[i] == "CLOSING" and (
                line_features.is_delimiter or line_features.is_visual_separator
            ):
                # Look ahead to determine context
                if i + 1 < len(fixed) and fixed[i + 1] == "SIGNATURE":
                    fixed[i] = "SIGNATURE"
                else:
                    fixed[i] = "BODY"
                logger.debug(
                    "Fixed separator line labeled CLOSING → %s at position %d",
                    fixed[i],
                    i,
                )

            if fixed[i] == "SIGNATURE":
                in_signature = True
            elif in_signature:
                # After entering SIGNATURE territory, very limited transitions allowed
                if fixed[i] == "CLOSING":
                    # SIGNATURE → CLOSING is impossible
                    fixed[i] = "SIGNATURE"
                    logger.debug("Fixed impossible SIGNATURE → CLOSING at position %d", i)

        return fixed

    def _unify_bracketed_blocks(
        self, labels: list[str], extracted: ExtractedFeatures
    ) -> list[str]:
        """Unify labels within bracketed blocks (separator-sandwiched sections).

        A bracketed block includes the opening separator, content lines, and
        closing separator. If the majority of lines in the block are BODY,
        the entire block becomes BODY. If the majority are SIGNATURE, the
        entire block becomes SIGNATURE. Other cases are left unchanged.

        This fixes cases where info blocks within signatures get mixed labels
        due to CRF transition preferences.
        """
        if len(labels) < 3:
            return labels

        fixed = labels.copy()

        # Find all separator positions
        separator_indices: list[int] = []
        for i, lf in enumerate(extracted.line_features):
            if lf.is_delimiter or lf.is_visual_separator:
                separator_indices.append(i)

        if len(separator_indices) < 2:
            return fixed

        # Track which separators have been used as closers
        used_as_closer: set[int] = set()

        # For each potential opening separator, find matching closer
        for i, open_idx in enumerate(separator_indices):
            if open_idx in used_as_closer:
                continue

            # Look for a matching closer among remaining separators
            for close_idx in separator_indices[i + 1:]:
                if close_idx in used_as_closer:
                    continue

                # Need at least 1 content line between separators
                if close_idx <= open_idx + 1:
                    continue

                # Found a potential block: [open_idx, close_idx] inclusive
                block_indices = list(range(open_idx, close_idx + 1))

                # Count BODY and SIGNATURE labels in the block
                body_count = sum(1 for idx in block_indices if fixed[idx] == "BODY")
                sig_count = sum(1 for idx in block_indices if fixed[idx] == "SIGNATURE")
                total = len(block_indices)

                # Majority rule: more than half must be BODY or SIGNATURE
                if body_count > total / 2:
                    # Unify to BODY
                    for idx in block_indices:
                        if fixed[idx] != "BODY":
                            logger.debug(
                                "Unified bracketed block: %s → BODY at position %d",
                                fixed[idx],
                                idx,
                            )
                            fixed[idx] = "BODY"
                    used_as_closer.add(close_idx)
                    break
                elif sig_count > total / 2:
                    # Unify to SIGNATURE
                    for idx in block_indices:
                        if fixed[idx] != "SIGNATURE":
                            logger.debug(
                                "Unified bracketed block: %s → SIGNATURE at position %d",
                                fixed[idx],
                                idx,
                            )
                            fixed[idx] = "SIGNATURE"
                    used_as_closer.add(close_idx)
                    break
                # else: no clear majority, try next potential closer

        return fixed

    def predict(
        self,
        extracted: ExtractedFeatures,
        texts: tuple[str, ...],
    ) -> SequenceLabelingResult:
        """Predict labels for email lines.

        Args:
            extracted: Feature extraction result from FeatureExtractor.
            texts: Original text lines corresponding to the features.

        Returns:
            SequenceLabelingResult with labeled lines and probabilities.

        Raises:
            RuntimeError: If no model is loaded.
        """
        if self._tagger is None:
            raise RuntimeError("No CRF model loaded. Call load_model() first.")

        if extracted.total_lines == 0:
            return SequenceLabelingResult(
                labeled_lines=(),
                sequence_probability=1.0,
            )

        # Extract features
        feature_seq = _extract_feature_sequence(extracted, texts)

        # Set the sequence for tagging
        self._tagger.set(feature_seq)

        # Get predicted labels
        predicted_labels = self._tagger.tag()

        # Fix impossible transitions
        predicted_labels = self._fix_impossible_transitions(predicted_labels, extracted)

        # Unify bracketed blocks (separator-sandwiched sections)
        predicted_labels = self._unify_bracketed_blocks(predicted_labels, extracted)

        # Get sequence probability (probability of the predicted sequence)
        sequence_prob = self._tagger.probability(predicted_labels)

        # Get marginal probabilities for each position
        model_labels = self._tagger.labels()
        labeled_lines: list[LabeledLine] = []

        for idx, (text, pred_label) in enumerate(zip(texts, predicted_labels, strict=True)):
            # Compute marginal probability for each label at this position
            label_probs: dict[Label, float] = {}
            for label in LABELS:
                if label in model_labels:
                    label_probs[label] = self._tagger.marginal(label, idx)
                else:
                    label_probs[label] = 0.0

            # Confidence is the marginal probability of the predicted label
            confidence = label_probs.get(pred_label, 0.0)

            # Validate label is one of the expected labels
            validated_label: Label
            if pred_label in LABELS:
                validated_label = pred_label
            else:
                logger.warning("Unknown label '%s' at position %d, defaulting to OTHER", pred_label, idx)
                validated_label = "OTHER"
                confidence = 0.0

            labeled_lines.append(
                LabeledLine(
                    text=text,
                    label=validated_label,
                    confidence=confidence,
                    label_probabilities=label_probs,
                )
            )

        return SequenceLabelingResult(
            labeled_lines=tuple(labeled_lines),
            sequence_probability=sequence_prob,
        )


class CRFTrainer:
    """Trainer for CRF sequence labeling models.

    Wraps python-crfsuite's Trainer with a convenient interface
    for training on labeled email data.
    """

    def __init__(
        self,
        algorithm: str = "lbfgs",
        c1: float = 0.1,
        c2: float = 0.1,
        max_iterations: int = 100,
        all_possible_transitions: bool = True,
    ) -> None:
        """Initialize the CRF trainer.

        Args:
            algorithm: Training algorithm. Options: 'lbfgs', 'l2sgd', 'ap', 'pa', 'arow'.
            c1: L1 regularization coefficient.
            c2: L2 regularization coefficient.
            max_iterations: Maximum number of training iterations.
            all_possible_transitions: Include transitions not in training data.
        """
        self._trainer = pycrfsuite.Trainer(verbose=False)
        self._trainer.set_params(
            {
                "c1": c1,
                "c2": c2,
                "max_iterations": max_iterations,
                "feature.possible_transitions": all_possible_transitions,
            }
        )
        self._trainer.select(algorithm)

    def add_sequence(
        self,
        extracted: ExtractedFeatures,
        texts: tuple[str, ...],
        labels: tuple[Label, ...],
    ) -> None:
        """Add a training sequence.

        Args:
            extracted: Feature extraction result from FeatureExtractor.
            texts: Original text lines.
            labels: Ground truth labels for each line.

        Raises:
            ValueError: If the number of labels doesn't match the number of lines.
        """
        if len(labels) != extracted.total_lines:
            raise ValueError(
                f"Number of labels ({len(labels)}) doesn't match "
                f"number of lines ({extracted.total_lines})"
            )

        if len(texts) != extracted.total_lines:
            raise ValueError(
                f"Number of texts ({len(texts)}) doesn't match "
                f"number of lines ({extracted.total_lines})"
            )

        feature_seq = _extract_feature_sequence(extracted, texts)
        self._trainer.append(feature_seq, list(labels))

    def train(self, output_path: Path | str) -> None:
        """Train the model and save to file.

        Args:
            output_path: Path to save the trained model.
        """
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)

        logger.info("Training CRF model...")
        self._trainer.train(str(path))
        logger.info("Saved CRF model to %s", path)

    def get_params(self) -> dict[str, str]:
        """Get current training parameters."""
        return dict(self._trainer.params())
