"""Event-time feature snapshots for the stock-specific impact model.

Features MUST use only pre-event data to prevent lookahead leakage.
Immutable snapshots are persisted at prediction time and never modified.

Design reference: Section I (Impact and Horizon Model) in design.md.
Requirement 12.2, 12.10.
"""

from __future__ import annotations

import hashlib
import json
from datetime import datetime

from pydantic import BaseModel, Field, field_validator

# ---------------------------------------------------------------------------
# Feature model
# ---------------------------------------------------------------------------


class ImpactFeatureSet(BaseModel):
    """Complete feature set for impact prediction.

    All features represent pre-event state. Timing rules:
    - Market features (volatility, volume, regime) use data strictly before event_time.
    - Extraction features (event class, sentiment, etc.) use the extraction output.
    - Company attributes use the most recent known state before event_time.

    Missing-value policy:
    - Numeric fields: NaN (float('nan')) when unavailable.
    - Categorical fields: "unknown" when unavailable.
    """

    # --- Event features (from v3 extraction) ---
    event_class_probabilities: dict[str, float] = Field(
        description="Probability distribution over event classes from specialist extractor.",
    )
    sentiment_positive: float = Field(
        description="Calibrated positive sentiment probability.",
    )
    sentiment_negative: float = Field(
        description="Calibrated negative sentiment probability.",
    )
    sentiment_neutral: float = Field(
        description="Calibrated neutral sentiment probability.",
    )
    magnitude: float = Field(
        description="Numeric magnitude/surprise of the event (NaN if unavailable).",
    )
    surprise: float = Field(
        description="Normalized surprise vs consensus or prior (NaN if unavailable).",
    )

    # --- Source features ---
    source_credibility: float = Field(
        description="Historical source accuracy score (NaN if unknown source).",
    )
    novelty_score: float = Field(
        description="Retrieval-based novelty score (0=duplicate, 1=completely novel).",
    )
    evidence_coverage: float = Field(
        description="Fraction of extracted facts backed by valid evidence spans.",
    )

    # --- Company attributes (pre-event snapshot) ---
    company_sector: str = Field(
        default="unknown",
        description="GICS sector or 'unknown'.",
    )
    company_industry: str = Field(
        default="unknown",
        description="GICS industry or 'unknown'.",
    )
    market_cap_bucket: str = Field(
        default="unknown",
        description="Market cap bucket: mega, large, mid, small, micro, unknown.",
    )
    beta: float = Field(
        description="Company beta relative to benchmark (NaN if unavailable).",
    )

    # --- Market state features (pre-event) ---
    pre_event_volatility: float = Field(
        description="Realized volatility in the lookback window before event (NaN if unavailable).",
    )
    volume_regime: str = Field(
        default="unknown",
        description="Volume regime: high, normal, low, unknown.",
    )
    broad_market_regime: str = Field(
        default="unknown",
        description="Broad market regime: bull, bear, choppy, unknown.",
    )

    # --- Event characterization ---
    event_directness: str = Field(
        default="unknown",
        description="Whether the event is direct, second_order, confirmed, quoted, speculative, or unknown.",
    )
    document_type: str = Field(
        default="unknown",
        description="Document type: news, filing, transcript, press_release, macro_event, unknown.",
    )

    # --- Metadata (not used as model inputs but needed for auditing) ---
    event_time: datetime = Field(
        description="Timestamp when the event was detected/published.",
    )
    feature_version: str = Field(
        default="1.0.0",
        description="Version of the feature extraction code.",
    )

    @field_validator("market_cap_bucket")
    @classmethod
    def validate_market_cap_bucket(cls, v: str) -> str:
        valid = {"mega", "large", "mid", "small", "micro", "unknown"}
        if v not in valid:
            return "unknown"
        return v

    @field_validator("volume_regime")
    @classmethod
    def validate_volume_regime(cls, v: str) -> str:
        valid = {"high", "normal", "low", "unknown"}
        if v not in valid:
            return "unknown"
        return v

    @field_validator("broad_market_regime")
    @classmethod
    def validate_broad_market_regime(cls, v: str) -> str:
        valid = {"bull", "bear", "choppy", "unknown"}
        if v not in valid:
            return "unknown"
        return v

    @field_validator("event_directness")
    @classmethod
    def validate_event_directness(cls, v: str) -> str:
        valid = {"direct", "second_order", "confirmed", "quoted", "speculative", "unknown"}
        if v not in valid:
            return "unknown"
        return v

    @field_validator("document_type")
    @classmethod
    def validate_document_type(cls, v: str) -> str:
        valid = {"news", "filing", "transcript", "press_release", "macro_event", "unknown"}
        if v not in valid:
            return "unknown"
        return v

    def to_numeric_vector(self) -> list[float]:
        """Convert to a flat numeric vector for tabular model input.

        Categorical fields are encoded as ordinal indices.
        NaN values are preserved for the model to handle (e.g., via missing-value support).
        """
        # Encode categoricals
        sector_map = {
            "Technology": 0, "Consumer Cyclical": 1, "Financial Services": 2,
            "Healthcare": 3, "Energy": 4, "Communication Services": 5,
            "Industrials": 6, "Consumer Defensive": 7, "Real Estate": 8,
            "Utilities": 9, "unknown": 10,
        }
        cap_map = {"mega": 0, "large": 1, "mid": 2, "small": 3, "micro": 4, "unknown": 5}
        volume_map = {"high": 0, "normal": 1, "low": 2, "unknown": 3}
        regime_map = {"bull": 0, "bear": 1, "choppy": 2, "unknown": 3}
        directness_map = {
            "direct": 0, "second_order": 1, "confirmed": 2,
            "quoted": 3, "speculative": 4, "unknown": 5,
        }
        doc_type_map = {
            "news": 0, "filing": 1, "transcript": 2,
            "press_release": 3, "macro_event": 4, "unknown": 5,
        }

        # Event class probabilities sorted by key for consistency
        event_probs = [
            self.event_class_probabilities.get(k, 0.0)
            for k in sorted(self.event_class_probabilities.keys())
        ] if self.event_class_probabilities else [0.0]

        return [
            *event_probs,
            self.sentiment_positive,
            self.sentiment_negative,
            self.sentiment_neutral,
            self.magnitude,
            self.surprise,
            self.source_credibility,
            self.novelty_score,
            self.evidence_coverage,
            float(sector_map.get(self.company_sector, 10)),
            float(cap_map.get(self.market_cap_bucket, 5)),
            self.beta,
            self.pre_event_volatility,
            float(volume_map.get(self.volume_regime, 3)),
            float(regime_map.get(self.broad_market_regime, 3)),
            float(directness_map.get(self.event_directness, 5)),
            float(doc_type_map.get(self.document_type, 5)),
        ]


# ---------------------------------------------------------------------------
# Feature snapshot persistence
# ---------------------------------------------------------------------------

# In-memory store for immutable snapshots (production would use object storage)
_FEATURE_SNAPSHOTS: dict[str, dict] = {}


def persist_feature_snapshot(features: ImpactFeatureSet, prediction_time: datetime) -> str:
    """Persist an immutable feature snapshot at prediction time.

    The snapshot is content-addressed: identical features at the same prediction
    time produce the same snapshot ID. Once written, snapshots are never modified.

    Parameters
    ----------
    features
        The complete feature set at event time.
    prediction_time
        When the prediction is being made (must be >= event_time).

    Returns
    -------
    str
        A unique, deterministic snapshot ID.

    Raises
    ------
    ValueError
        If prediction_time is before the feature event_time (temporal inconsistency).
    """
    if prediction_time < features.event_time:
        raise ValueError(
            f"prediction_time ({prediction_time.isoformat()}) cannot be before "
            f"event_time ({features.event_time.isoformat()})"
        )

    # Serialize deterministically for content-addressing
    snapshot_data = {
        "features": features.model_dump(mode="json"),
        "prediction_time": prediction_time.isoformat(),
    }
    content = json.dumps(snapshot_data, sort_keys=True, default=str)
    snapshot_id = hashlib.sha256(content.encode()).hexdigest()[:16]

    # Immutable write — never overwrite
    if snapshot_id not in _FEATURE_SNAPSHOTS:
        _FEATURE_SNAPSHOTS[snapshot_id] = snapshot_data

    return snapshot_id


def get_feature_snapshot(snapshot_id: str) -> dict | None:
    """Retrieve a persisted feature snapshot by ID."""
    return _FEATURE_SNAPSHOTS.get(snapshot_id)


def clear_feature_snapshots() -> None:
    """Clear all stored snapshots (for testing only)."""
    _FEATURE_SNAPSHOTS.clear()


# ---------------------------------------------------------------------------
# Timing validation
# ---------------------------------------------------------------------------


def validate_no_future_leakage(
    features: ImpactFeatureSet,
    market_data_timestamps: list[datetime] | None = None,
) -> list[str]:
    """Check that no feature uses post-event data.

    Parameters
    ----------
    features
        The feature set to validate.
    market_data_timestamps
        Optional list of timestamps from market data used in features.
        All must be strictly before event_time.

    Returns
    -------
    list[str]
        List of leakage violations found (empty = no leakage).
    """
    violations: list[str] = []
    event_time = features.event_time

    if market_data_timestamps:
        for i, ts in enumerate(market_data_timestamps):
            if ts >= event_time:
                violations.append(
                    f"market_data_timestamps[{i}] ({ts.isoformat()}) is at or after "
                    f"event_time ({event_time.isoformat()})"
                )

    return violations