from collections import deque
import tensorflow as tf

from fsm_crawl.blocking.feature_extractor import FeatureExtractor
from fsm_crawl.blocking.model import BlockingModel


class BlockingController:
    def __init__(
        self,
        model: BlockingModel = BlockingModel(),
        feature_extractor: FeatureExtractor = FeatureExtractor(),
        block_requests: bool = True,
        train_online: bool = True,
        batch_size: int = 512,
        threshold: float = 0.5,
    ):
        self.model = model
        self.fe = feature_extractor
        self.block_requests = block_requests
        self.train_online = train_online
        self.batch_size = batch_size
        self.threshold = threshold

        self.X_buf = deque(maxlen=batch_size)
        self.y_buf = deque(maxlen=batch_size)

    async def handle_request(self, route, request):
        # Feature extraction + label
        features, label = self.fe.extract(request)  # features: [203], label: 0/1
        features = tf.expand_dims(features, axis=0)  # [1, 203]

        # Predict
        prob = self.model.predict(features)[0, 0]

        # Blocking
        if self.block_requests and prob > self.threshold:
            await route.abort()
        else:
            await route.continue_()

        # Buffer for training
        self.X_buf.append(features)
        self.y_buf.append(label)
        print(len(self.X_buf), self.batch_size)
        if self.train_online and len(self.X_buf) == self.batch_size:
            X_batch, y_batch = self._make_balanced_batch()
            self.model.train_batch(X_batch, y_batch)
            self.X_buf.clear()
            self.y_buf.clear()

    def _make_balanced_batch(self):
        """
        Vectorized balanced sampling (fast).
        """
        X = tf.concat(list(self.X_buf), axis=0)         # [B, 203]
        y = tf.convert_to_tensor(self.y_buf, tf.int32)  # [B]

        # Split by class
        pos_idx = tf.where(y == 1)[:, 0]
        neg_idx = tf.where(y == 0)[:, 0]

        half = self.batch_size // 2

        # Guard against empty classes
        if tf.size(pos_idx) == 0 or tf.size(neg_idx) == 0:
            return X, tf.cast(y, tf.float32)

        # Resample indices (with replacement)
        pos_sampled = tf.gather(
            pos_idx,
            tf.random.uniform([half], 0, tf.shape(pos_idx)[0], dtype=tf.int32),
        )
        neg_sampled = tf.gather(
            neg_idx,
            tf.random.uniform([half], 0, tf.shape(neg_idx)[0], dtype=tf.int32),
        )

        idx = tf.concat([pos_sampled, neg_sampled], axis=0)

        X_balanced = tf.gather(X, idx)
        y_balanced = tf.concat(
            [tf.ones(half), tf.zeros(half)], axis=0
        )

        return X_balanced, tf.cast(y_balanced, tf.float32)
