# -*- coding: utf-8 -*-
"""ELM Web file loader class."""
import asyncio
import logging
from pathlib import Path
from abc import ABC, abstractmethod

import aiohttp
from fake_useragent import UserAgent

from elm.utilities.parse import read_pdf
from elm.web.document import PDFDocument, HTMLDocument
from elm.web.html_pw import load_html_with_pw
from elm.web.utilities import DEFAULT_HEADERS
from elm.utilities.retry import async_retry_with_exponential_backoff
from elm.exceptions import ELMRuntimeError


logger = logging.getLogger(__name__)


async def _read_pdf_doc(pdf_bytes, **kwargs):
    """Default read PDF function (runs in main thread)"""
    verbose = kwargs.pop("verbose", True)
    pages = read_pdf(pdf_bytes, verbose=verbose)
    return PDFDocument(pages, **kwargs)


async def _read_html_doc(text, **kwargs):
    """Default read HTML function (runs in main thread)"""
    return HTMLDocument([text], **kwargs)


async def _read_pdf_file(pdf_fp, **kwargs):
    """Default read PDF file function (runs in main thread)"""
    verbose = kwargs.pop("verbose", True)
    with open(pdf_fp, "rb") as fh:
        pdf_bytes = fh.read()
    pages = read_pdf(pdf_bytes, verbose=verbose)
    return PDFDocument(pages, **kwargs), pdf_bytes


async def _read_html_file(html_fp, **kwargs):
    """Default read HTML function (runs in main thread)"""
    with open(html_fp, "r") as fh:
        text = fh.read()
    return HTMLDocument([text], **kwargs), text


class BaseAsyncFileLoader(ABC):
    """Base class for async file loading"""

    def __init__(
        self,
        pdf_read_coroutine,
        html_read_coroutine,
        pdf_read_kwargs=None,
        html_read_kwargs=None,
        pdf_ocr_read_coroutine=None,
        file_cache_coroutine=None,
        **__,  # consume any extra kwargs
    ):
        """

        Parameters
        ----------
        pdf_read_coroutine : callable
            PDF file read coroutine. Must by an async function.
            Must return a :obj:`elm.web.document.PDFDocument`.
        html_read_coroutine : callable, optional
            HTML file read coroutine. Must by an async function.
            Must return a :obj:`elm.web.document.HTMLDocument`.
        pdf_read_kwargs : dict, optional
            Keyword-value argument pairs to pass to the
            `pdf_read_coroutine`. By default, ``None``.
        html_read_kwargs : dict, optional
            Keyword-value argument pairs to pass to the
            `html_read_coroutine`. By default, ``None``.
        pdf_ocr_read_coroutine : callable, optional
            PDF OCR file read coroutine. Must by an async function.
            Should accept PDF bytes as the first argument and kwargs as
            the rest. Must return a :obj:`elm.web.document.PDFDocument`.
            If ``None``, PDF OCR parsing is not attempted, and any
            scanned PDF URL's will return a blank document.
            By default, ``None``.
        file_cache_coroutine : callable, optional
            File caching coroutine. Can be used to cache files
            downloaded by this class. Must accept an
            :obj:`~elm.web.document.Document` instance as the first
            argument and the file content to be written as the second
            argument. If this method is not provided, no document
            caching is performed. By default, ``None``.
        """
        self.pdf_read_coroutine = pdf_read_coroutine
        self.html_read_coroutine = html_read_coroutine
        self.pdf_read_kwargs = pdf_read_kwargs or {}
        self.html_read_kwargs = html_read_kwargs or {}
        self.pdf_ocr_read_coroutine = pdf_ocr_read_coroutine
        self.file_cache_coroutine = file_cache_coroutine

    async def fetch_all(self, *sources):
        """Fetch documents for all requested sources.

        Parameters
        ----------
        *sources
            Iterable of sources (as strings) used to fetch the
            documents.

        Returns
        -------
        list
            List of documents, one per requested sources.
        """
        outer_task_name = asyncio.current_task().get_name()
        fetches = [
            asyncio.create_task(self.fetch(source), name=outer_task_name)
            for source in sources
        ]
        return await asyncio.gather(*fetches)

    async def fetch(self, source):
        """Fetch a document for the given source.

        Parameters
        ----------
        source : str
            Source used to load the document.

        Returns
        -------
        :class:`elm.web.document.Document`
            Document instance containing text, if the load was
            successful.
        """
        try:
            doc, raw = await self._fetch_doc_with_url_in_metadata(source)
        except KeyboardInterrupt:
            raise
        except Exception as e:
            msg = ("Encountered error of type %r while fetching document from "
                   "%s:")
            err_type = type(e)
            logger.exception(msg, err_type, source)
            return HTMLDocument(pages=[])

        doc = await self._cache_doc(doc, raw)
        return doc

    async def _fetch_doc_with_url_in_metadata(self, source):
        """Fetch doc contents and add source to metadata"""
        doc, raw_content = await self._fetch_doc(source)
        doc.attrs["source"] = source
        return doc, raw_content

    async def _cache_doc(self, doc, raw_content):
        """Cache doc if user provided a coroutine"""
        if doc.empty or not raw_content:
            return doc

        if not self.file_cache_coroutine:
            return doc

        cache_fn = await self.file_cache_coroutine(doc, raw_content)
        if cache_fn is not None:
            doc.attrs["cache_fn"] = cache_fn
        return doc

    @abstractmethod
    async def _fetch_doc(self, source):
        """Fetch documents given a source"""
        raise NotImplementedError


class AsyncWebFileLoader(BaseAsyncFileLoader):
    """Async web file (PDF or HTML) loader

    Purpose:
        Save content from links as files.
    Responsibilities:
        1. Retrieve data from a URL.
        2. Determine wether information should be stored as a PDF or
           HTML document.
    Key Relationships:
        Returns either :class:`~elm.web.document.PDFDocument` or
        :class:`~elm.web.document.HTMLDocument`. Uses `aiohttp` to
        access the web.

    .. end desc
    """

    PAGE_LOAD_TIMEOUT = 60_000
    """Default page load timeout value in milliseconds"""

    def __init__(
        self,
        header_template=None,
        verify_ssl=True,
        aget_kwargs=None,
        pw_launch_kwargs=None,
        pdf_read_kwargs=None,
        html_read_kwargs=None,
        pdf_read_coroutine=None,
        html_read_coroutine=None,
        pdf_ocr_read_coroutine=None,
        file_cache_coroutine=None,
        browser_semaphore=None,
        use_scrapling_stealth=False,
        num_pw_html_retries=3,
        **__,  # consume any extra kwargs
    ):
        """

        Parameters
        ----------
        header_template : dict, optional
            Optional GET header template. If not specified, uses
            :obj:`~elm.web.utilities.DEFAULT_HEADERS`.
            By default, ``None``.
        verify_ssl : bool, optional
            Option to use aiohttp's default SSL check. If ``False``,
            SSL certificate validation is skipped. By default, ``True``.
        aget_kwargs : dict, optional
            Other kwargs to pass to :meth:`aiohttp.ClientSession.get`.
            By default, ``None``.
        pw_launch_kwargs : dict, optional
            Keyword-value argument pairs to pass to
            :meth:`async_playwright.chromium.launch` (only used when
            reading HTML). By default, ``None``.
        pdf_read_kwargs : dict, optional
            Keyword-value argument pairs to pass to the
            `pdf_read_coroutine`. By default, ``None``.
        html_read_kwargs : dict, optional
            Keyword-value argument pairs to pass to the
            `html_read_coroutine`. By default, ``None``.
        pdf_read_coroutine : callable, optional
            PDF file read coroutine. Must by an async function. Should
            accept PDF bytes as the first argument and kwargs as the
            rest. Must return a :obj:`elm.web.document.PDFDocument`.
            If ``None``, a default function that runs in the main thread
            is used. By default, ``None``.
        html_read_coroutine : callable, optional
            HTML file read coroutine. Must by an async function. Should
            accept HTML text as the first argument and kwargs as the
            rest. Must return a :obj:`elm.web.document.HTMLDocument`.
            If ``None``, a default function that runs in the main thread
            is used. By default, ``None``.
        pdf_ocr_read_coroutine : callable, optional
            PDF OCR file read coroutine. Must by an async function.
            Should accept PDF bytes as the first argument and kwargs as
            the rest. Must return a :obj:`elm.web.document.PDFDocument`.
            If ``None``, PDF OCR parsing is not attempted, and any
            scanned PDF URL's will return a blank document.
            By default, ``None``.
        file_cache_coroutine : callable, optional
            File caching coroutine. Can be used to cache files
            downloaded by this class. Must accept an
            :obj:`~elm.web.document.Document` instance as the first
            argument and the file content to be written as the second
            argument. If this method is not provided, no document
            caching is performed. By default, ``None``.
        browser_semaphore : asyncio.Semaphore, optional
            Semaphore instance that can be used to limit the number of
            playwright browsers open concurrently. If ``None``, no
            limits are applied. By default, ``None``.
        use_scrapling_stealth : bool, default=False
            Option to use scrapling stealth scripts instead of
            tf-playwright-stealth. By default, ``False``.
        num_pw_html_retries : int, default=3
            Number of attempts to load HTML content. This is useful
            because the playwright parameters are stochastic, and
            sometimes a combination of them can fail to load HTML. The
            default value is likely a good balance between processing
            attempts and retrieval success. Note that the minimum number
            of attempts will always be 2, even if the user provides a
            value smaller than this. By default, ``3``.
        """
        super().__init__(
            pdf_read_coroutine=pdf_read_coroutine or _read_pdf_doc,
            html_read_coroutine=html_read_coroutine or _read_html_doc,
            pdf_read_kwargs=pdf_read_kwargs,
            html_read_kwargs=html_read_kwargs,
            pdf_ocr_read_coroutine=pdf_ocr_read_coroutine,
            file_cache_coroutine=file_cache_coroutine
        )
        self.pw_launch_kwargs = pw_launch_kwargs or {}
        self.get_kwargs = {
            "headers": self._header_from_template(header_template),
            "ssl": None if verify_ssl else False,
            **(aget_kwargs or {}),
        }
        self.browser_semaphore = browser_semaphore
        self.uss = use_scrapling_stealth
        self.num_pw_html_retries = num_pw_html_retries

    def _header_from_template(self, header_template):
        """Compile header from user or default template"""
        headers = header_template or DEFAULT_HEADERS
        headers = dict(headers)
        if not headers.get("User-Agent"):
            headers["User-Agent"] = UserAgent().random
        return headers

    async def _fetch_doc(self, url):
        """Fetch a doc by trying pdf read, then HTML read, then PDF OCR"""

        async with aiohttp.ClientSession() as session:
            try:
                logger.debug("Fetching content from %r", url)
                out = await self._fetch_content_with_retry(url, session)
            except ELMRuntimeError:
                logger.exception("Could not fetch content from %r", url)
                return PDFDocument(pages=[]), None

        raw_content, ct, charset = out
        logger.debug("Got content from %r", url)
        doc = await self.pdf_read_coroutine(raw_content,
                                            **self.pdf_read_kwargs)
        if not doc.empty:
            return doc, raw_content

        logger.debug("PDF read failed; fetching HTML content from %r", url)
        doc = await self._fetch_html_using_pw_with_retry(url)
        if not doc.empty:
            return doc, doc.text

        if "text" in ct:
            logger.debug("HTML read with playwright failed; fetching HTML "
                         "content from response with content type %r and "
                         "charset %r for %r", ct, charset, url)
            doc = await self._try_load_doc_from_response_text(raw_content,
                                                              charset)
            if not doc.empty:
                return doc, doc.text

        elif self.pdf_ocr_read_coroutine:
            logger.debug("HTML read failed; fetching OCR content from %r", url)
            doc = await self.pdf_ocr_read_coroutine(
                raw_content, **self.pdf_read_kwargs
            )

        return doc, raw_content

    async def _fetch_html_using_pw_with_retry(self, url):
        """Fetch HTML content with several retry attempts"""
        num_attempts = max(1, int(self.num_pw_html_retries) - 1)
        max_attempts = num_attempts + 1
        for attempt in range(num_attempts):
            logger.debug("HTML read for %r (attempt %d of %d)",
                         url, attempt + 1, max_attempts)
            text = await load_html_with_pw(url, self.browser_semaphore,
                                           timeout=self.PAGE_LOAD_TIMEOUT,
                                           use_scrapling_stealth=self.uss,
                                           **self.pw_launch_kwargs)
            doc = await self.html_read_coroutine(text, **self.html_read_kwargs)
            if not doc.empty:
                return doc

        logger.debug("HTML read for %r (attempt %d of %d) with "
                     "load_state='domcontentloaded'",
                     url, max_attempts, max_attempts)
        text = await load_html_with_pw(url, self.browser_semaphore,
                                       timeout=self.PAGE_LOAD_TIMEOUT,
                                       use_scrapling_stealth=self.uss,
                                       load_state="domcontentloaded",
                                       **self.pw_launch_kwargs)
        return await self.html_read_coroutine(text, **self.html_read_kwargs)

    @async_retry_with_exponential_backoff(
        base_delay=2,
        exponential_base=1.5,
        jitter=False,
        max_retries=3,
        errors=(
            aiohttp.ClientConnectionError,
            aiohttp.client_exceptions.ClientError,
        ),
    )
    async def _fetch_content_with_retry(self, url, session):
        """Fetch content from URL with several retry attempts"""
        async with session.get(url, **self.get_kwargs) as response:
            body = await response.read()
            ct = response.content_type.casefold()
            charset = response.charset or 'utf-8'
            return body, ct, charset

    async def _try_load_doc_from_response_text(self, raw_content, charset):
        """Try to load document by decoding response text"""
        try:
            text = raw_content.decode(charset)
        except Exception:
            return HTMLDocument(pages=[])

        return await self.html_read_coroutine(text, **self.html_read_kwargs)


class AsyncLocalFileLoader(BaseAsyncFileLoader):
    """Async local file (PDF or HTML) loader"""

    def __init__(
        self,
        pdf_read_kwargs=None,
        html_read_kwargs=None,
        pdf_read_coroutine=None,
        html_read_coroutine=None,
        pdf_ocr_read_coroutine=None,
        file_cache_coroutine=None,
        doc_attrs=None,
        **__,  # consume any extra kwargs
    ):
        """

        Parameters
        ----------
        pdf_read_kwargs : dict, optional
            Keyword-value argument pairs to pass to the
            `pdf_read_coroutine`. By default, ``None``.
        html_read_kwargs : dict, optional
            Keyword-value argument pairs to pass to the
            `html_read_coroutine`. By default, ``None``.
        pdf_read_coroutine : callable, optional
            PDF file read coroutine. Must by an async function. Should
            accept a PDF filepath as the first argument and kwargs as
            the rest. Must return a :obj:`elm.web.document.PDFDocument`
            along with the raw PDF bytes (for caching purposes).
            If ``None``, a default function that runs in the main thread
            is used. By default, ``None``.
        html_read_coroutine : callable, optional
            HTML file read coroutine. Must by an async function. Should
            accept an HTML filepath as the first argument and kwargs as
            the rest. Must return a :obj:`elm.web.document.HTMLDocument`
            along with the raw text (for caching purposes).
            If ``None``, a default function that runs in the main thread
            is used. By default, ``None``.
        pdf_ocr_read_coroutine : callable, optional
            PDF OCR file read coroutine. Must by an async function.
            Should accept a PDF filepath as the first argument and
            kwargs as the rest. Must return a
            :obj:`elm.web.document.PDFDocument` along with the raw PDF
            bytes (for caching purposes).
            If ``None``, PDF OCR parsing is not attempted, and any
            scanned PDF URL's will return a blank document.
            By default, ``None``.
        file_cache_coroutine : callable, optional
            File caching coroutine. Can be used to cache files
            downloaded by this class. Must accept an
            :obj:`~elm.web.document.Document` instance as the first
            argument and the file content to be written as the second
            argument. If this method is not provided, no document
            caching is performed. By default, ``None``.
        doc_attrs : dict, optional
            Additional document attributes to add to each loaded
            document. By default, ``None``.
        """
        super().__init__(
            pdf_read_coroutine=pdf_read_coroutine or _read_pdf_file,
            html_read_coroutine=html_read_coroutine or _read_html_file,
            pdf_read_kwargs=pdf_read_kwargs,
            html_read_kwargs=html_read_kwargs,
            pdf_ocr_read_coroutine=pdf_ocr_read_coroutine,
            file_cache_coroutine=file_cache_coroutine
        )
        self.doc_attrs = doc_attrs or {}

    async def _fetch_doc(self, source):
        """Load a doc by reading file based on extension"""
        fp = Path(source)
        if fp.suffix.lower() == ".pdf":
            logger.debug("Trying to read PDF file: %r", source)
            doc, raw = await self.pdf_read_coroutine(fp,
                                                     **self.pdf_read_kwargs)
            if not doc.empty:
                return doc, raw
            elif self.pdf_ocr_read_coroutine:
                logger.debug("PDF read failed; fetching OCR content from %r",
                             source)
                doc, raw = await self.pdf_ocr_read_coroutine(
                    fp, **self.pdf_read_kwargs)
                if not doc.empty:
                    return doc, raw

        if fp.suffix.lower() == ".txt":
            logger.debug("Trying to read HTML file: %r", source)
            doc, raw = await self.html_read_coroutine(fp,
                                                      **self.html_read_kwargs)
            if not doc.empty:
                return doc, raw

        logger.error("Failed to read file: %r", source)
        return PDFDocument(pages=[]), None

    async def _fetch_doc_with_url_in_metadata(self, source):
        """Fetch doc contents and add source to metadata"""
        doc, raw_content = await self._fetch_doc(source)
        for key, value in self.doc_attrs.items():
            doc.attrs[key] = value
        doc.attrs["source_fp"] = source
        return doc, raw_content


class AsyncFileLoader(AsyncWebFileLoader):
    """Alias for AsyncWebFileLoader (for backward compatibility)"""
