import tensorflow as tf
from urllib.parse import urlparse
from fsm_crawl.blocking.labeler import Labeler
METHODS = ["GET", "POST", "OPTIONS", "HEAD", "PUT", "DELETE", "SEARCH", "PATCH"]
TYPES = [
    "xmlhttprequest", "image", "font", "script", "stylesheet", "ping",
    "sub_frame", "other", "main_frame", "csp_report", "object", "media"
]

class FeatureExtractor:
    URL_LENGTH = 200
    
    def __init__(self, labeler: Labeler = Labeler()):
        self.labeler = labeler

    def generate_url_encoding(self, url: str):
        encoding = [(ord(c) % 89) + 1 for c in url]
        if len(encoding) < self.URL_LENGTH:
            encoding = [0] * (self.URL_LENGTH - len(encoding)) + encoding
        else:
            encoding = encoding[-self.URL_LENGTH:]
        return encoding

    def extract(self, request) -> tuple[tf.Tensor, int]:
        features = self.generate_url_encoding(request.url)
        features.append(METHODS.index(request.method) + 1 if request.method in METHODS else 0)
        features.append(TYPES.index(request.resource_type) + 1 if request.resource_type in TYPES else 0)
        page_domain = urlparse(request.frame.url).hostname if request.frame else None
        req_domain = urlparse(request.url).hostname
        features.append(1 if page_domain != req_domain else 0)
        label = self.labeler.label(request.url)
        return tf.reshape(tf.convert_to_tensor(features, dtype=tf.float32), [203]), label
