# This script crawls a website and checks for broken links and size limits.
# The order of the links or the referrer is not guaranteed to be stable.

import requests
from bs4 import BeautifulSoup
from urllib.parse import urljoin, urlparse
import json
import sys
import logging
import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import os


def setup_logging(debug=False, output_file=None):
    """Configure logging with handlers for file (INFO) and stderr (ERROR and DEBUG)"""
    # Create our own logger
    logger = logging.getLogger('crawler')
    logger.handlers = []  # Clear any existing handlers
    logger.propagate = False  # Don't propagate to root logger

    # Set base level based on debug flag
    logger.setLevel(logging.DEBUG if debug else logging.INFO)

    # Create formatter
    formatter = logging.Formatter('%(message)s')

    # Create file handler for INFO level (JSONL output)
    if output_file:
        file_handler = logging.FileHandler(output_file, mode='w')
        file_handler.setLevel(logging.INFO)
        file_handler.setFormatter(formatter)
        file_handler.addFilter(lambda record: record.levelno == logging.INFO)
        logger.addHandler(file_handler)

    # Create stderr handler for ERROR and DEBUG
    stderr_handler = logging.StreamHandler(sys.stderr)
    stderr_handler.setLevel(logging.DEBUG)
    stderr_handler.setFormatter(formatter)
    stderr_handler.addFilter(lambda record: record.levelno != logging.INFO)
    logger.addHandler(stderr_handler)

    return logger


def get_domain(url):
    """Extracts the (scheme, netloc) part of a URL to identify its domain."""
    parsed = urlparse(url)
    return parsed.scheme, parsed.netloc


def normalize_url(url):
    """Normalize URL by removing fragment identifier."""
    parsed = urlparse(url)
    return urljoin(url, parsed.path + ('?' + parsed.query if parsed.query else ''))


def is_same_domain(url, base_domain):
    """Checks whether 'url' is in the same domain or a subdomain of 'base_domain'."""
    scheme, netloc = get_domain(url)
    base_scheme, base_netloc = base_domain

    # Extract the main domain from both URLs by taking the last two parts
    # e.g., 'sub.example.com' -> 'example.com'
    main_domain = '.'.join(netloc.split('.')[-2:])
    base_main_domain = '.'.join(base_netloc.split('.')[-2:])

    return scheme == base_scheme and main_domain == base_main_domain


def is_valid_url(url):
    """Check if the URL is valid and not a JavaScript pseudo-URL."""
    if url.startswith(('javascript:', 'void(', '#', 'tel:', 'mailto:')):
        return False
    try:
        result = urlparse(url)
        return all([result.scheme, result.netloc])
    except Exception:
        return False


def process_url(url_info):
    """Process a single URL and return its data."""
    url, base_domain, should_recurse, referrer, depth = url_info
    result = {"url": url, "links": [], "referrer": referrer}
    logger = logging.getLogger('crawler')

    if not is_valid_url(url):
        logger.debug(f"Skipping invalid URL: {url}")
        return url, {"status_code": "invalid", "size": 0, "links": [], "referrer": referrer}

    try:
        headers = {
            'User-Agent': 'curl/8.7.1',
            'Accept': '*/*'
        }
        response = requests.get(
            url, timeout=10, allow_redirects=True, headers=headers)
        status_code = response.status_code

        if 400 <= status_code < 600:
            logger.error(f"Found error status code {status_code}: {url}")
            return url, {"status_code": status_code, "size": 0, "links": [], "referrer": None}
        elif status_code == "error":
            logger.info(f"Found status code {status_code}: {url}")
            return url, {"status_code": 0, "size": 0, "links": [], "referrer": None}

        page_content = response.content
        size_in_bytes = len(page_content)

        # Check if the final URL after redirects is still in our domain
        final_url = response.url
        if not is_same_domain(final_url, base_domain):
            logger.debug(
                f"URL {url} redirected outside our domain to {final_url}")
            return url, {"status_code": status_code, "size": size_in_bytes, "links": [], "referrer": referrer}

        result.update({"status_code": status_code, "size": size_in_bytes})

        # Only proceed if it's HTML
        content_type = response.headers.get("Content-Type", "")
        if "text/html" in content_type.lower():
            soup = BeautifulSoup(page_content, "html.parser")
            for link_tag in soup.find_all("a", href=True):
                raw_link = link_tag.get("href")
                # Use final_url instead of original url
                next_link = urljoin(final_url, raw_link)
                normalized_link = normalize_url(next_link)

                # For internal links - continue crawling with no depth limit
                if is_same_domain(normalized_link, base_domain):
                    # Keep recursing for internal links if within depth limit
                    result["links"].append(
                        (normalized_link, True, url, depth + 1))
                # For external links - only check them once (no recursion)
                elif should_recurse:  # Only add external links from internal pages
                    # External links never recurse
                    result["links"].append(
                        (normalized_link, False, url, depth + 1))

            logger.debug(f"Found {len(result['links'])} new links from {url}")

    except requests.RequestException as e:
        logger.error(f"Error fetching {url}: {e}")
        return url, {"status_code": "error", "size": 0, "links": [], "referrer": None}

    return url, result


def crawl_site(start_url, max_workers=10, max_depth=None):
    logger = logging.getLogger('crawler')
    logger.debug(f"Starting crawl from {start_url}")
    if not start_url.startswith(('http://', 'https://')):
        start_url = 'http://' + start_url

    base_domain = get_domain(start_url)
    logger.debug(f"Base domain: {base_domain[0]}://{base_domain[1]}")

    # Track URLs to process and those we've seen, along with their original referrers
    normalized_start = normalize_url(start_url)
    # None referrer for start URL, True for should_recurse since it's the starting point, depth 0 for start
    to_process = [(normalized_start, base_domain, True, None, 0)]
    seen_urls = {normalized_start: None}  # Track referrers in seen_urls
    results = {}

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        while to_process:
            # Process URLs in parallel batches
            future_to_url = {
                executor.submit(process_url, url_info): url_info[0]
                for url_info in to_process
            }

            to_process = []

            for future in as_completed(future_to_url):
                url = future_to_url[future]
                try:
                    url, data = future.result()
                    # Set the referrer from our tracking, not from the URL
                    data["referrer"] = seen_urls[url]
                    results[url] = data

                    # Add new links to process if within depth limit
                    for new_url, new_should_recurse, _, new_depth in data.get("links", []):
                        if new_url not in seen_urls and (max_depth is None or new_depth <= max_depth):
                            # Store the referrer when we first see the URL
                            seen_urls[new_url] = url
                            to_process.append(
                                (new_url, base_domain, new_should_recurse, url, new_depth))

                except Exception as e:
                    logger.error(f"Error processing {url}: {e}")
                    results[url] = {"status_code": "error",
                                    "size": 0, "referrer": seen_urls.get(url)}

    logger.debug(f"Crawl complete. Visited {len(results)} URLs")
    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Crawl a website and check for broken links and size limits')
    parser.add_argument('url', help='The URL to start crawling from')
    parser.add_argument('--max-size', type=int,
                        help='Maximum allowed size in bytes for any page', default=None)
    parser.add_argument('--workers', type=int,
                        help='Number of parallel workers', default=10)
    parser.add_argument('--max-depth', type=int,
                        help='Maximum depth to crawl (default: no limit)', default=None)
    parser.add_argument('--verbose', action='store_true',
                        help='Enable verbose logging')
    args = parser.parse_args()

    # Get domain name for the output file
    if not args.url.startswith(('http://', 'https://')):
        args.url = 'http://' + args.url
    domain = urlparse(args.url).netloc
    output_file = f"result_{domain}.jsonl"

    # Setup logging with the output file
    logger = setup_logging(args.verbose, output_file)

    report = crawl_site(args.url, max_workers=args.workers,
                        max_depth=args.max_depth)

    has_errors = False
    size_exceeded = False
    final_errors = []
    # Sort URLs by size before output
    sorted_items = sorted(report.items(), key=lambda x: x[1]["size"])
    # Output in JSONL format
    for url, info in sorted_items:
        result = {
            "url": url,
            "status_code": info["status_code"],
            "size": info["size"],
            "referrer": info.get("referrer")
        }
        logger.info(json.dumps(result))
        if isinstance(info["status_code"], int) and 400 <= info["status_code"] < 600:
            final_errors.append(
                f"Error: {url} returned status code {info['status_code']}")
            has_errors = True
        if args.max_size and info["size"] > args.max_size and is_same_domain(url, args.url):
            final_errors.append(
                f"Error: {url} exceeds maximum size of {args.max_size} bytes (actual size: {info['size']} bytes)")
            size_exceeded = True

    # Print all collected errors at the end
    if final_errors:
        for error in final_errors:
            logger.error(error)

    # Print the path to results file
    abs_output_path = os.path.abspath(output_file)
    print(f"\nResults written to: {abs_output_path}")

    sys.exit(has_errors or size_exceeded)
