from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Iterable, Optional
from datetime import date

import polars as pl

from ocf.pipelines.steps import (
    step_normalize_raw,
    step_align_underlying,
    step_enrich_chain,
    step_per_option_greeks,
    step_features_underlying,
    assemble_pipeline_outputs,
)
from ocf.schemas.validation import validate_canonical_bundle


# Pipeline configuration
@dataclass
class PipelineConfig:
    """
    Configuration for the options-chain feature pipeline.
    """

    # core
    symbol: str
    snapshot_date: date

    # feature selection
    include_features: Optional[Iterable[str]] = None
    exclude_features: Optional[Iterable[str]] = None

    # risk-free rate selection
    rate_field: str = "rate_sofr_90d"

    # toggles
    build_per_option: bool = True

    # validation / safety
    require_iv_surface: bool = False


# Internal shared execution
def _run_from_normalized(
    *,
    config: PipelineConfig,
    normalized: Dict[str, pl.DataFrame],
) -> Dict[str, pl.DataFrame]:
    """
    Execute pipeline steps starting from normalized inputs.
    """

    # Align underlying
    underlying_daily = step_align_underlying(
        normalized=normalized
    )

    # Enrich option chain snapshot
    chain_enriched = step_enrich_chain(
        chain=normalized["chain"],
        underlying_daily=underlying_daily,
    )

    # Per-option Greeks
    if config.build_per_option:
        per_option = step_per_option_greeks(
            chain_enriched=chain_enriched,
            underlying_daily=underlying_daily,
            rate_field=config.rate_field,
        )
    else:
        per_option = pl.DataFrame([])

    # Underlying features
    features = step_features_underlying(
        underlying_daily=underlying_daily,
        include=config.include_features,
        exclude=config.exclude_features,
    )

    # Assemble canonical outputs
    outputs = assemble_pipeline_outputs(
        underlying_daily=underlying_daily,
        features=features,
        chain_snapshot=chain_enriched,
        per_option_greeks=per_option,
    )

    # Validate canonical outputs
    validate_canonical_bundle(
        outputs,
        require_chain=config.build_per_option,
        require_iv_surface=config.require_iv_surface,
    )

    return outputs


# Pipeline runner
class FeaturePipeline:
    """
    Orchestrates the full options-chain feature pipeline.
    """

    def __init__(self, config: PipelineConfig):
        self.config = config

    def run(
        self,
        *,
        raw_inputs: Dict[str, str | pl.DataFrame],
    ) -> Dict[str, pl.DataFrame]:
        """
        Executes the full pipeline starting from raw vendor inputs.
        """

        # Normalize
        normalized = step_normalize_raw(
            symbol=self.config.symbol,
            raw=raw_inputs,
            snapshot_date=self.config.snapshot_date,
        )

        return _run_from_normalized(
            config=self.config,
            normalized=normalized,
        )


# Public APIs
def run_pipeline(
    *,
    config: PipelineConfig,
    raw_inputs: Dict[str, str | pl.DataFrame],
) -> Dict[str, pl.DataFrame]:
    """
    Functional entry point starting from raw inputs.
    """
    pipeline = FeaturePipeline(config)
    return pipeline.run(raw_inputs=raw_inputs)


def run_pipeline_normalized(
    *,
    config: PipelineConfig,
    normalized_inputs: Dict[str, pl.DataFrame],
) -> Dict[str, pl.DataFrame]:
    """
    Functional entry point starting from pre-normalized inputs.

    The caller is responsible for ensuring canonical schema correctness.
    """
    return _run_from_normalized(
        config=config,
        normalized=normalized_inputs,
    )