"""
Streaming data processor with contract-driven pipelines.

This module provides the core streaming processor that:
1. Reads streaming contracts (YAML)
2. Automatically selects framework (Bytewax vs Pathway)
3. Creates streaming pipelines
4. Validates data quality in real-time
5. Writes to Bronze (history) and Realtime (current) layers
"""

from typing import TYPE_CHECKING, Any, Dict, Iterator

from loguru import logger

if TYPE_CHECKING:
    from bytewax.dataflow import Dataflow


class StreamingFrameworkSelector:
    """
    Automatically select the best streaming framework based on contract.
    """

    def select(self, contract: Dict[str, Any]) -> str:
        """
        Select streaming framework based on contract analysis.

        Args:
            contract: Parsed contract dictionary

        Returns:
            "bytewax" or "pathway"
        """
        # Check if user explicitly specified framework
        streaming_config = contract.get("source", {}).get("streaming", {})
        framework = streaming_config.get("framework", "auto")

        if framework and framework != "auto":
            logger.info(f"Using user-specified framework: {framework}")
            return framework

        # Auto-select based on contract features
        logger.info("Auto-selecting streaming framework...")

        score_bytewax = 0
        score_pathway = 0
        reasons = []

        # Rule 1: SQL-heavy operations → Pathway
        if self._has_complex_sql(contract):
            score_pathway += 3
            reasons.append("Complex SQL operations detected → Pathway")

        # Rule 2: Stream-batch joins → Pathway
        if self._has_stream_batch_joins(contract):
            score_pathway += 5
            reasons.append("Stream-batch joins detected → Pathway")

        # Rule 3: Complex aggregations → Pathway
        if self._has_complex_aggregations(contract):
            score_pathway += 3
            reasons.append("Complex aggregations detected → Pathway")

        # Rule 4: High throughput → Bytewax
        if self._requires_high_throughput(contract):
            score_bytewax += 3
            reasons.append("High throughput required → Bytewax")

        # Rule 5: Stateful processing → Bytewax
        if self._has_stateful_operations(contract):
            score_bytewax += 2
            reasons.append("Stateful operations detected → Bytewax")

        # Rule 6: Simple transformations → Bytewax
        if self._has_simple_transformations(contract):
            score_bytewax += 1
            reasons.append("Simple transformations → Bytewax")

        # Default: Bytewax (if scores are equal)
        if score_pathway > score_bytewax:
            selected = "pathway"
        else:
            selected = "bytewax"

        logger.info(f"✅ Auto-selected framework: {selected}")
        logger.info(f"   Bytewax score: {score_bytewax}, Pathway score: {score_pathway}")
        for reason in reasons:
            logger.info(f"   - {reason}")

        return selected

    def _has_complex_sql(self, contract: Dict) -> bool:
        """Check if contract has complex SQL operations."""
        # Check realtime aggregations for SQL functions
        realtime = contract.get("realtime", [])
        for agg in realtime:
            for metric in agg.get("metrics", []):
                expr = metric.get("expression", "").upper()
                # Complex SQL functions
                if any(func in expr for func in ["CASE", "COALESCE", "NULLIF", "MODE", "PERCENTILE"]):
                    return True

        # Check transformations for SQL
        transformations = contract.get("transformations", [])
        for transform in transformations:
            expr = transform.get("expression", "").upper()
            if any(func in expr for func in ["CASE", "COALESCE", "NULLIF"]):
                return True

        return False

    def _has_stream_batch_joins(self, contract: Dict) -> bool:
        """Check if contract has stream-batch joins."""
        enrichment = contract.get("enrichment", [])
        for enrich in enrichment:
            if enrich.get("type") == "join" and enrich.get("source") in [
                "delta",
                "parquet",
                "csv",
            ]:
                return True
        return False

    def _has_complex_aggregations(self, contract: Dict) -> bool:
        """Check if contract has complex aggregations."""
        realtime = contract.get("realtime", [])

        # Multiple aggregations → complex
        if len(realtime) > 2:
            return True

        # Multiple group-by columns → complex
        for agg in realtime:
            if len(agg.get("group_by", [])) > 2:
                return True

            # Multiple metrics → complex
            if len(agg.get("metrics", [])) > 3:
                return True

        return False

    def _requires_high_throughput(self, contract: Dict) -> bool:
        """Check if high throughput is required."""
        streaming = contract.get("source", {}).get("streaming", {})
        throughput = streaming.get("throughput", {})

        expected_rps = throughput.get("expected_records_per_second", 0)

        # > 10K records/sec → high throughput
        return expected_rps > 10000

    def _has_stateful_operations(self, contract: Dict) -> bool:
        """Check if contract has stateful operations."""
        streaming = contract.get("source", {}).get("streaming", {})
        window = streaming.get("window", {})

        # Session windows → stateful
        if window.get("type") == "session":
            return True

        # Deduplication → stateful
        stream_rules = contract.get("quality", {}).get("stream_rules", [])
        for rule in stream_rules:
            if rule.get("rule") == "deduplication":
                return True

        return False

    def _has_simple_transformations(self, contract: Dict) -> bool:
        """Check if transformations are simple."""
        transformations = contract.get("transformations", [])

        # No transformations → simple
        if not transformations:
            return True

        # All transformations are simple expressions
        for transform in transformations:
            expr = transform.get("expression", "").upper()
            # Complex if has CASE, subqueries, etc.
            if any(keyword in expr for keyword in ["CASE", "SELECT", "JOIN"]):
                return False

        return True


class StreamingDataProcessor:
    """
    Process streaming data with contracts.

    Features:
    - Contract-driven pipeline generation
    - Automatic framework selection (Bytewax vs Pathway)
    - Real-time data validation
    - Bronze (history) + Realtime (current) dual-write
    - Automatic retention management

    Example:
        >>> processor = StreamingDataProcessor(
        ...     contract="contracts/wikimedia_stream.yaml"
        ... )
        >>> processor.start()  # Runs continuously
    """

    def __init__(self, contract: str, framework: str = "auto"):
        """
        Initialize streaming processor.

        Args:
            contract: Path to contract YAML file
            framework: "auto", "bytewax", or "pathway"
        """
        self.contract_path = contract
        self.contract = self._load_contract(contract)

        # Select framework
        selector = StreamingFrameworkSelector()
        self.framework = selector.select(self.contract) if framework == "auto" else framework

        logger.info(f"Initialized StreamingDataProcessor with {self.framework}")

    def _load_contract(self, contract_path: str) -> Dict[str, Any]:
        """Load contract from YAML file."""
        import yaml

        with open(contract_path, "r") as f:
            contract = yaml.safe_load(f)

        logger.info(f"Loaded contract: {contract.get('dataset', 'unknown')}")
        return contract

    def start(self):
        """
        Start streaming pipeline.

        This runs continuously until stopped (Ctrl+C).
        """
        logger.info(f"Starting streaming pipeline with {self.framework}...")

        if self.framework == "bytewax":
            self._start_bytewax()
        elif self.framework == "pathway":
            self._start_pathway()
        else:
            raise ValueError(f"Unknown framework: {self.framework}")

    def _start_bytewax(self):
        """Start Bytewax streaming pipeline."""
        try:
            from bytewax import operators as op
            from bytewax.connectors.stdio import StdOutSink
            from bytewax.dataflow import Dataflow
        except ImportError:
            raise ImportError("Bytewax not installed. Install with: pip install 'lakelogic[bytewax]'")

        logger.info("Creating Bytewax dataflow...")

        # Create dataflow
        flow = Dataflow("lakelogic_streaming")

        # 1. Input from source
        stream = op.input("source", flow, self._create_input())

        # 2. Validate schema
        stream = op.map("validate_schema", stream, self._validate_schema)

        # 3. Apply transformations
        stream = op.map("transform", stream, self._apply_transformations)

        # 4. Validate quality
        stream = op.map("validate_quality", stream, self._validate_quality)

        # 5. Split good/bad records
        good_stream = op.filter("filter_good", stream, lambda x: x.get("_quality_status") == "good")
        _bad_stream = op.filter("filter_bad", stream, lambda x: x.get("_quality_status") == "bad")

        # 6. Write to Bronze (history)
        op.output("bronze", good_stream, StdOutSink())  # TODO: Replace with Delta Lake

        # 7. Create realtime aggregations
        for agg_config in self.contract.get("realtime", []):
            self._create_bytewax_aggregation(flow, good_stream, agg_config)

        # Run dataflow
        logger.info("✅ Bytewax dataflow created, starting execution...")
        logger.info("Press Ctrl+C to stop")

        # For now, just run with stdio (will add proper execution later)
        from bytewax.testing import run_main

        run_main(flow)

    def _start_pathway(self):
        """Start Pathway streaming pipeline."""
        try:
            import pathway as pw
        except ImportError:
            raise ImportError("Pathway not installed. Install with: pip install 'lakelogic[pathway]'")

        logger.info("Creating Pathway pipeline...")

        # 1. Define input table from the contract source
        # We pass the stream iterator directly to Pathway
        table = pw.debug.table_from_iterable(self._create_input())

        # 2. Schema Validation (UDF)
        @pw.udf
        def validate_schema_udf(record: Dict[str, Any]) -> Dict[str, Any]:
            return self._validate_schema(record)

        table = table.select(data=validate_schema_udf(pw.this))

        # 3. Transformations (UDF)
        @pw.udf
        def transform_udf(record: Dict[str, Any]) -> Dict[str, Any]:
            return self._apply_transformations(record)

        table = table.select(data=transform_udf(pw.this.data))

        # 4. Quality Validation (UDF)
        @pw.udf
        def validate_quality_udf(record: Dict[str, Any]) -> Dict[str, Any]:
            return self._validate_quality(record)

        table = table.select(data=validate_quality_udf(pw.this.data))

        # 5. Split Good/Bad and Sink
        good_table = table.filter(pw.this.data["_quality_status"] == "good")
        bad_table = table.filter(pw.this.data["_quality_status"] == "bad")

        # 6. Sinks (For now, print to console using debug sink)
        pw.io.csv.write(good_table, "quarantine/good_stream.csv")
        pw.io.csv.write(bad_table, "quarantine/bad_stream.csv")

        # Execution
        logger.info("✅ Pathway pipeline created, starting execution...")
        pw.run()

    def _create_input(self) -> Iterator[Dict[str, Any]]:
        """Create input stream from contract source."""
        source_config = self.contract.get("source", {})
        source_type = source_config.get("type")
        connector = source_config.get("connector")

        if source_type == "stream":
            if connector == "sse":
                # SSE connector (Wikimedia, etc.)
                from lakelogic.engines.streaming_connectors import SSEConnector

                url = source_config.get("url")
                conn = SSEConnector(url)

                logger.info(f"Connected to SSE stream: {url}")

                for event in conn.stream():
                    yield event

            elif connector == "websocket":
                # WebSocket connector (Coinbase, etc.)
                from lakelogic.engines.streaming_connectors import WebSocketConnector

                url = source_config.get("url")
                subscribe_message = source_config.get("subscribe_message")

                conn = WebSocketConnector(url, subscribe_message)

                logger.info(f"Connected to WebSocket: {url}")

                for event in conn.stream():
                    yield event

            elif connector == "kafka":
                # Kafka connector
                from lakelogic.engines.streaming_connectors import KafkaConnector

                brokers = source_config.get("brokers", ["localhost:9092"])
                topic = source_config.get("topic")
                group = source_config.get("consumer_group", "lakelogic")

                conn = KafkaConnector(brokers, topic, group)
                for event in conn.stream():
                    yield event

            elif connector == "webhook":
                # Webhook connector (Receiver)
                from lakelogic.engines.streaming_connectors import WebhookConnector

                port = source_config.get("port", 8080)
                path = source_config.get("path", "/webhook")

                conn = WebhookConnector(port=port, path=path)
                for event in conn.stream():
                    yield event

            else:
                raise ValueError(f"Unknown connector: {connector}")

        else:
            raise ValueError(f"Unknown source type: {source_type}")

    def _validate_schema(self, event: Dict[str, Any]) -> Dict[str, Any]:
        """Validate event schema against contract model."""
        model = self.contract.get("model", {})
        _fields = model.get("fields", [])

        # For now, just pass through (full validation in next update)
        # TODO: Implement schema validation

        return event

    def _apply_transformations(self, event: Dict[str, Any]) -> Dict[str, Any]:
        """Apply transformations from contract."""
        _transformations = self.contract.get("transformations", [])

        # For now, just pass through (full transformations in next update)
        # TODO: Implement transformations

        return event

    def _validate_quality(self, event: Dict[str, Any]) -> Dict[str, Any]:
        """Validate data quality rules from contract."""
        quality = self.contract.get("quality", {})
        _row_rules = quality.get("row_rules", [])

        # For now, mark all as good (full validation in next update)
        # TODO: Implement quality validation
        event["_quality_status"] = "good"

        return event

    def _create_bytewax_aggregation(self, flow: "Dataflow", stream: Any, agg_config: Dict[str, Any]):
        """Create Bytewax aggregation from contract config."""
        # TODO: Implement aggregations
        # For now, just log
        logger.info(f"Aggregation configured: {agg_config.get('name')}")
        pass


def select_streaming_framework(contract: Dict[str, Any]) -> str:
    """
    Public API to select streaming framework.

    Args:
        contract: Parsed contract dictionary

    Returns:
        "bytewax" or "pathway"

    Example:
        >>> import yaml
        >>> with open("contract.yaml") as f:
        ...     contract = yaml.safe_load(f)
        >>> framework = select_streaming_framework(contract)
        >>> print(framework)
        'bytewax'
    """
    selector = StreamingFrameworkSelector()
    return selector.select(contract)
