# Utility Functions
import logging
import os
import re
import asyncio
import httpx
import validators
from typing import Optional
from urllib.parse import urlparse
from bs4 import BeautifulSoup

logger = logging.getLogger(__name__)

class InvalidURLException(Exception):
    """Exception raised for invalid URLs."""
    pass

class InsecureURLException(Exception):
    """Exception raised for insecure URLs."""
    pass

class SusURLException(Exception):
    """Exception raised when VirusTotal flags a URL as unsafe or the URL is unreachable."""
    pass

class FetchingError(Exception):
    """Exception raised for errors during fetching."""
    pass

async def check_url_safety(url: str, max_retries=5, retry_delay=5) -> Optional[bool]:
    """
    Check URL safety by submitting to VirusTotal and polling until analysis completes.

    Submits the URL to the VirusTotal scan endpoint, then polls the analysis endpoint
    until the status is 'completed' or retries are exhausted.

    Args:
        url: The URL to analyze (used in error logging only; analysis_id from VirusTotal is the actual lookup key).
        max_retries: Maximum number of polling attempts while the analysis is queued.
        retry_delay: Seconds to wait between polling attempts.

    Returns:
        True if the URL is safe (0 malicious detections, fewer than 3 suspicious detections).
        False if flagged as unsafe (malicious > 0 OR suspicious >= 3). None if the analysis
        could not complete (API errors, unexpected status, or missing data in response).
    """
    headers = {'x-apikey': os.environ['TELLMGRAMBOT_VIRUSTOTAL_API_KEY']}
    data = {'url': url}

    async with httpx.AsyncClient() as client:
        response = await client.post(url='https://www.virustotal.com/api/v3/urls', headers=headers, data=data)

        if response.status_code == 200:  # OK
            json_response = response.json()
            analysis_id = json_response.get("data", {}).get("id", None)  # Extract the analysis ID from the response
            if analysis_id:
                retries = 0  # Initialize retry counter
                while retries <= max_retries:  # While we haven't exceeded max retries
                    # Now check the analysis results
                    response = await client.get(
                        url=f"https://www.virustotal.com/api/v3/analyses/{analysis_id}",
                        headers=headers
                    )
                    if response.status_code == 200:  # OK
                        json_response = response.json()
                        attributes = json_response.get("data", {}).get("attributes", None)
                        if attributes:
                            status = attributes.get("status", None)
                            if status == "completed":
                                stats = attributes.get("stats")
                                if stats is None:
                                    logger.warning("Stats not found in VirusTotal response.")
                                    return None
                                malicious = stats.get("malicious", 0)
                                suspicious = stats.get("suspicious", 0)
                                if malicious == 0 and suspicious < 3:
                                    return True  # URL is safe
                                else:
                                    logger.warning(
                                        f"URL is not safe. Malicious: {malicious}, "
                                        f"Suspicious: {suspicious}, Status: {status}"
                                    )
                                    return False  # URL is not safe
                            elif status == "queued":
                                retries += 1  # Increment retry counter
                                await asyncio.sleep(retry_delay)  # Wait before retrying
                            else:
                                logger.warning(f"Unexpected VirusTotal status: {status}")
                                return None
                        else:
                            logger.warning("Attributes not found in VirusTotal response.")
                            return None
                    else:
                        logger.warning("Error fetching VirusTotal analysis results.")
                        return None
            else:
                logger.warning("Analysis ID not found in VirusTotal response.")
                return None
        else:
            logger.warning(f"VirusTotal request failed with status: {response.status_code}")
            return None

async def fetch_url(url: str, div_id: Optional[str] = None, div_class: Optional[str] = None) -> Optional[str]:
    """
    Validate, safety-check, and fetch the full or scoped content of a URL.

    Performs URL validation, HTTPS enforcement, and VirusTotal safety check. If validation passes,
    fetches the full HTML content and optionally returns only a scoped <div> tag by id or class.

    Args:
        url: The URL to fetch. Must use HTTPS protocol.
        div_id: Optional <div> id attribute to extract only that element's HTML content.
        div_class: Optional <div> class attribute to extract only the first matching element's HTML content.
                   Ignored if div_id is specified.

    Returns:
        The full raw HTML if neither div_id nor div_class is specified. If div_id or
        div_class is provided, returns the HTML of the first matching <div>, or None
        if no matching <div> is found.
    """
    try:
        async with httpx.AsyncClient() as client:
            # Validate the URL
            if not validators.url(url):
                raise InvalidURLException(f"[{url}] is invalid!")
            elif not urlparse(url).scheme.startswith('https'):
                raise InsecureURLException(f"[{url}] is insecure!")
            safety = await check_url_safety(url)
            if safety is False:
                raise SusURLException(f"[{url}] is neither safe nor existent!")
            elif safety is None:
                raise FetchingError(f"VirusTotal analysis failed for [{url}], cannot verify safety.")

            # Fetch the URL if everything is OK
            response = await client.get(url)
            response.raise_for_status()  # Raises an exception for 4xx and 5xx responses

            # Parse the HTML content
            soup = BeautifulSoup(response.text, features='html.parser')

            # If a specific <div> tag ID or class is provided, extract only that content
            specific_div = None
            if div_id:
                specific_div = soup.find(name='div', id=div_id)
            elif div_class:
                specific_div = soup.find(name='div', class_=div_class)

            if specific_div:
                return str(specific_div)  # Returns the HTML of the specified <div> tag
            elif div_id or div_class:
                logger.info(f"Specified <div> tag with ID {div_id} or class {div_class} not found")
                return None  # Returns none but they were expecting to find something
            else:
                return response.text  # Returns the entire HTML content as they didn't specify a <div> tag

    except httpx.HTTPError as err:
        raise FetchingError(f"An unhandled HTTP error occurred while fetching [{url}] contents: {err}")

def strip_html_markup(html_content: str) -> str:
    """
    Extract plain text from an HTML string, collapsing blank lines.

    Parses the HTML content using BeautifulSoup and returns only the text content
    with all tags removed and runs of blank lines collapsed to a single newline.

    Args:
        html_content: Raw HTML string to parse.

    Returns:
        Plain text with all HTML tags removed, consecutive newlines collapsed to single
        newlines, and leading/trailing whitespace stripped.
    """
    soup = BeautifulSoup(html_content, features='html.parser')
    text_content = soup.get_text()
    cleaned_text = re.sub('\n+', '\n', text_content).strip()
    return cleaned_text
