"""Predefined tools for LLMs"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_tools.ipynb.

# %% auto 0
__all__ = ['scrape_url', 'read_url']

# %% ../nbs/04_tools.ipynb 2
from bs4 import BeautifulSoup
from cloudscraper import create_scraper
from html2text import HTML2Text
from urllib.parse import urljoin, urlparse
import re
from textwrap import dedent

# %% ../nbs/04_tools.ipynb 5
def _absolutify_imgs(md, base_url):
    """This function rewrites Markdown image links so their URLs become absolute, using a base URL.
    - md: a Markdown string
    - base_url: the base URL used to resolve relative path
    """
    def fix(m):
        alt, img_url = m.group(1), m.group(2)
        if not img_url.startswith('http'): img_url = urljoin(base_url, img_url)
        return f'![{alt}]({img_url})'
    return re.sub(r'!\[(.*?)\]\((.*?)\)', fix, md)


# %% ../nbs/04_tools.ipynb 6
def _convert_math(soup, mode):
    """This function walks an HTML/XML document, finds MathML math elements, and replaces them with TeX/LaTeX markup suitable for text-based math renderers (like Markdown, MathJax, or KaTeX).
    - soup: a BeautifulSoup object representing parsed HTML/XML
    - mode: controls the inline math delimiter style
    Two inline math styles are supported:
    - mode='dollar': Inline math $x$ Display math $$x$$
    - other	Inline math \\(...\\) Display math $$x$$
    """
    for math in soup.find_all('math'):
        annot = math.find('annotation', {'encoding': 'application/x-tex'})
        if not annot:
            continue
        tex, display = annot.text.strip(), math.get('display') == 'block'
        if mode == 'dollar':
            wrap = f'$${tex}$$' if display else f'${tex}$'
        else:
            wrap = f'$${tex}$$' if display else f'\\({tex}\\)'
        math.replace_with(wrap)

# %% ../nbs/04_tools.ipynb 7
def scrape_url(url):
    "Get the html content of a web page using the cloudscraper library to bypass Cloudflare's anti-bot page."
    return create_scraper().get(url)

def read_url(url: str, as_md: bool = True, extract_section: bool = True, selector: str = None, math_mode: str = None):
    """This functions extracts a web page information for LLM ingestion
    1. Downloads a web page
    2. Parses HTML
    3. Optionally extracts a specific section (fragment or CSS selector)
    4. Converts MathML → LaTeX
    5. Optionally converts HTML → Markdown
    6. Convert code sections to fenced markdown blocks
    7. Makes image URLs absolute
    8. Returns the processed text
    """
    o = scrape_url(url)
    res, ctype = o.text, o.headers.get('content-type').split(';')[0]
    soup = BeautifulSoup(res, 'lxml')

    if selector:
        res = '\n\n'.join(str(s) for s in soup.select(selector))
    elif extract_section:
        parsed = urlparse(url)
        if parsed.fragment:
            section = soup.find(id=parsed.fragment)
            if section:
                elements = [section]
                current = section.next_sibling
                while current:
                    if hasattr(current, 'name') and current.name == section.name: break
                    elements.append(current)
                    current = current.next_sibling
                res = ''.join(str(el) for el in elements)
            else:
                res = ''
    else:
        res = str(soup)

    if math_mode:
        res_soup = BeautifulSoup(res, 'lxml')
        _convert_math(res_soup, math_mode)
        res = str(res_soup)

    if as_md and ctype == 'text/html':
        h = HTML2Text()
        h.body_width = 0
        # Handle code blocks
        h.mark_code = True
        res = h.handle(res)
        def _f(m): return f'```\n{dedent(m.group(1))}\n```'
        res = re.sub(r'\[code]\s*\n(.*?)\n\[/code]', _f, res or '', flags=re.DOTALL).strip()
        # Handle image urls
        res = _absolutify_imgs(res, urljoin(url, s['href'] if (s := soup.find('base')) else ''))
        # Handle math blocks
        if math_mode == 'safe':
            res = res.replace('\\\\(', '\\(').replace('\\\\)', '\\)')

    return res

# %% ../nbs/04_tools.ipynb 25
# import ...
