import gzip
import xml.etree.ElementTree as ET

import httpx

from utils.logger import logger


class SitemapParser:
    """
    Parse sitemap.xml files
    to discover documentation URLs.

    Supports:
    - sitemap.xml
    - sitemap index
    - gzipped sitemap
    """

    def __init__(self, timeout=30):
        self.timeout = timeout

    async def parse(
        self,
        sitemap_url: str
    ):
        """
        Main entry point

        Returns:
        -------
        List[str]
        """

        logger.info(
            f"Parsing sitemap: {sitemap_url}"
        )

        content = await self.fetch_sitemap(
            sitemap_url
        )

        if not content:
            return []

        root = ET.fromstring(content)

        tag = self.clean_tag(root.tag)

        # sitemap index
        if tag == "sitemapindex":

            urls = await self.parse_sitemap_index(
                root
            )

            logger.info(
                f"Sitemap index found {len(urls)} URLs"
            )

            return urls

        # regular sitemap
        elif tag == "urlset":

            urls = self.parse_urlset(root)

            logger.info(
                f"Sitemap contains {len(urls)} URLs"
            )

            return urls

        else:

            logger.warning(
                f"Unknown sitemap type: {tag}"
            )

            return []

    async def fetch_sitemap(
        self,
        sitemap_url: str
    ):
        """
        Download sitemap content
        """

        try:

            async with httpx.AsyncClient(
                timeout=self.timeout,
                follow_redirects=True
            ) as client:

                response = await client.get(
                    sitemap_url
                )

                response.raise_for_status()

                content = response.content

                # gzip sitemap
                if sitemap_url.endswith(".gz"):

                    content = gzip.decompress(
                        content
                    )

                return content

        except Exception as e:

            logger.exception(
                f"Failed fetching sitemap {sitemap_url}: {e}"
            )

            return None

    async def parse_sitemap_index(
        self,
        root
    ):
        """
        Parse sitemap index recursively

        Example:

        <sitemapindex>
            <sitemap>
                <loc>...</loc>
            </sitemap>
        </sitemapindex>
        """

        all_urls = []

        for sitemap in root.findall(".//*"):

            tag = self.clean_tag(sitemap.tag)

            if tag != "loc":
                continue

            child_sitemap_url = sitemap.text

            if not child_sitemap_url:
                continue

            logger.info(
                f"Found child sitemap: {child_sitemap_url}"
            )

            try:

                urls = await self.parse(
                    child_sitemap_url
                )

                all_urls.extend(urls)

            except Exception as e:

                logger.warning(
                    f"Failed child sitemap "
                    f"{child_sitemap_url}: {e}"
                )

        return all_urls

    def parse_urlset(
        self,
        root
    ):
        """
        Parse regular sitemap URL list

        Example:

        <urlset>
            <url>
                <loc>...</loc>
            </url>
        </urlset>
        """

        urls = []

        for url_node in root.findall(".//*"):

            tag = self.clean_tag(url_node.tag)

            if tag != "loc":
                continue

            url = url_node.text

            if not url:
                continue

            url = url.strip()

            if url not in urls:
                urls.append(url)

        return urls

    def clean_tag(
        self,
        tag: str
    ):
        """
        Remove XML namespace

        Example:
        {http://www.sitemaps.org/schemas/sitemap/0.9}url

        ->
        url
        """

        if "}" in tag:
            return tag.split("}", 1)[1]

        return tag

    async def discover_common_sitemaps(
        self,
        base_url: str
    ):
        """
        Try common sitemap locations

        Returns:
        -------
        List[str]
        """

        common_paths = [
            "/sitemap.xml",
            "/sitemap_index.xml",
            "/sitemap-index.xml",
            "/sitemap.gz",
            "/robots.txt"
        ]

        discovered = []

        for path in common_paths:

            url = base_url.rstrip("/") + path

            try:

                async with httpx.AsyncClient(
                    timeout=self.timeout
                ) as client:

                    response = await client.get(url)

                    if response.status_code == 200:

                        logger.info(
                            f"Found sitemap resource: {url}"
                        )

                        discovered.append(url)

            except Exception:
                pass

        return discovered

    async def parse_robots_for_sitemaps(
        self,
        robots_url: str
    ):
        """
        Extract sitemap URLs from robots.txt

        Example:

        Sitemap: https://example.com/sitemap.xml
        """

        sitemap_urls = []

        try:

            async with httpx.AsyncClient(
                timeout=self.timeout
            ) as client:

                response = await client.get(
                    robots_url
                )

                response.raise_for_status()

                text = response.text

            for line in text.splitlines():

                line = line.strip()

                if line.lower().startswith(
                    "sitemap:"
                ):

                    sitemap_url = line.split(
                        ":",
                        1
                    )[1].strip()

                    sitemap_urls.append(
                        sitemap_url
                    )

            logger.info(
                f"Found {len(sitemap_urls)} "
                f"sitemaps in robots.txt"
            )

        except Exception as e:

            logger.warning(
                f"Failed parsing robots.txt "
                f"{robots_url}: {e}"
            )

        return sitemap_urls