import logging
import re
from typing import List

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TextClassificationPipeline,
)

from guardrail.firewall.util import device

from .base_detector import Detector

_model_path = "elftsdmr/malware-url-detect"

log = logging.getLogger(__name__)

# Regular expression pattern to match URLs
url_pattern = re.compile(
    r"http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+"
)

class MalwareOutputURL(Detector):
    """
    MalwareURL Class:

    This class scans and detects malicious URLs in text using a pretrained model.
    It safeguards users from potential threats by classifying URLs as either malicious or benign.
    """

    def __init__(self, threshold=0.75):
        """
        Initializes an instance of the MalwareURL class.

        Parameters:
            threshold (float): The threshold used to determine if a URL is malicious. Defaults to 0.75.
        """

        model = AutoModelForSequenceClassification.from_pretrained(_model_path)
        self._tokenizer = AutoTokenizer.from_pretrained(_model_path)
        self._threshold = threshold
        self._text_classification_pipeline = TextClassificationPipeline(
            model=model,
            tokenizer=self._tokenizer,
            device=device,
        )
        log.debug(f"Initialized model {_model_path} on device {device}")

    @staticmethod
    def extract_urls(text: str) -> List[str]:
        """
        Extracts URLs from a given text.

        Args:
            text (str): The text to extract URLs from.

        Returns:
            List[str]: A list of URLs found in the text.
        """
        return url_pattern.findall(text)

    def scan(self, prompt: str, output: str) -> (str, bool, float):
        if output.strip() == "":
            return output, True, 0.0

        urls = self.extract_urls(output)
        if len(urls) == 0:
            return output, True, 0.0

        log.debug(f"Found {len(urls)} URLs in the output")

        urls_str = ", ".join(urls)
        result = self._text_classification_pipeline(
            urls_str, truncation=True, padding=True, max_length=self._tokenizer.model_max_length
        )
        malware_score = (
            result[0]["score"] if result[0]["label"] == "MALWARE" else 1 - result[0]["score"]
        )
        if malware_score > self._threshold:
            log.warning(
                f"Detected a malware URL with score: {malware_score}, threshold: {self._threshold}"
            )

            return output, False, round(malware_score, 2)

        log.debug(
            f"No malware URLs detected in the output. Max score: {malware_score}, threshold: {self._threshold}"
        )

        return output, True, 0.0
