# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
#
# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
# SPDX-License-Identifier: MIT
import glob
import hashlib
import logging
import os
import re
from typing import Any, Callable, Optional, Union
from urllib.parse import urlparse

import requests

from .import_utils import optional_import_block, require_optional_import
from .token_count_utils import count_token

with optional_import_block():
    import chromadb
    import markdownify
    from bs4 import BeautifulSoup

    if chromadb.__version__ < "0.4.15":
        from chromadb.api import API
    else:
        from chromadb.api import ClientAPI as API  # noqa: N814
    import chromadb.utils.embedding_functions as ef
    import pypdf
    from chromadb.api.types import QueryResult


with optional_import_block() as result:
    from unstructured.partition.auto import partition

HAS_UNSTRUCTURED = result.is_successful

logger = logging.getLogger(__name__)
TEXT_FORMATS = [
    "txt",
    "json",
    "csv",
    "tsv",
    "md",
    "html",
    "htm",
    "rtf",
    "rst",
    "jsonl",
    "log",
    "xml",
    "yaml",
    "yml",
    "pdf",
    "mdx",
]
UNSTRUCTURED_FORMATS = [
    "doc",
    "docx",
    "epub",
    "msg",
    "odt",
    "org",
    "pdf",
    "ppt",
    "pptx",
    "rtf",
    "rst",
    "xlsx",
]  # These formats will be parsed by the 'unstructured' library, if installed.
if HAS_UNSTRUCTURED:
    TEXT_FORMATS += UNSTRUCTURED_FORMATS
    TEXT_FORMATS = list(set(TEXT_FORMATS))
VALID_CHUNK_MODES = frozenset({"one_line", "multi_lines"})
RAG_MINIMUM_MESSAGE_LENGTH = int(os.environ.get("RAG_MINIMUM_MESSAGE_LENGTH", 5))


def split_text_to_chunks(
    text: str,
    max_tokens: int = 4000,
    chunk_mode: str = "multi_lines",
    must_break_at_empty_line: bool = True,
    overlap: int = 0,  # number of overlapping lines
):
    """Split a long text into chunks of max_tokens."""
    if chunk_mode not in VALID_CHUNK_MODES:
        raise AssertionError
    if chunk_mode == "one_line":
        must_break_at_empty_line = False
        overlap = 0
    chunks = []
    lines = text.split("\n")
    num_lines = len(lines)
    if num_lines < 3 and must_break_at_empty_line:
        logger.warning("The input text has less than 3 lines. Set `must_break_at_empty_line` to `False`")
        must_break_at_empty_line = False
    lines_tokens = [count_token(line) for line in lines]
    sum_tokens = sum(lines_tokens)
    while sum_tokens > max_tokens:
        estimated_line_cut = 2 if chunk_mode == "one_line" else max(int(max_tokens / sum_tokens * len(lines)), 2)
        cnt = 0
        prev = ""
        for cnt in reversed(range(estimated_line_cut)):
            if must_break_at_empty_line and lines[cnt].strip() != "":
                continue
            if sum(lines_tokens[:cnt]) <= max_tokens:
                prev = "\n".join(lines[:cnt])
                break
        if cnt == 0:
            logger.warning(
                f"max_tokens is too small to fit a single line of text. Breaking this line:\n\t{lines[0][:100]} ..."
            )
            if not must_break_at_empty_line:
                split_len = max(
                    int(max_tokens / (lines_tokens[0] * 0.9 * len(lines[0]) + 0.1)), RAG_MINIMUM_MESSAGE_LENGTH
                )
                prev = lines[0][:split_len]
                lines[0] = lines[0][split_len:]
                lines_tokens[0] = count_token(lines[0])
            else:
                logger.warning("Failed to split docs with must_break_at_empty_line being True, set to False.")
                must_break_at_empty_line = False
        (
            chunks.append(prev) if len(prev) >= RAG_MINIMUM_MESSAGE_LENGTH else None
        )  # don't add chunks less than RAG_MINIMUM_MESSAGE_LENGTH characters
        lines = lines[cnt - overlap if cnt > overlap else cnt :]
        lines_tokens = lines_tokens[cnt - overlap if cnt > overlap else cnt :]
        sum_tokens = sum(lines_tokens)
    text_to_chunk = "\n".join(lines).strip()
    (
        chunks.append(text_to_chunk) if len(text_to_chunk) >= RAG_MINIMUM_MESSAGE_LENGTH else None
    )  # don't add chunks less than RAG_MINIMUM_MESSAGE_LENGTH characters
    return chunks


@require_optional_import("pypdf", "retrievechat")
def extract_text_from_pdf(file: str) -> str:
    """Extract text from PDF files"""
    text = ""
    with open(file, "rb") as f:
        reader = pypdf.PdfReader(f)
        if reader.is_encrypted:  # Check if the PDF is encrypted
            try:
                reader.decrypt("")
            except pypdf.errors.FileNotDecryptedError as e:
                logger.warning(f"Could not decrypt PDF {file}, {e}")
                return text  # Return empty text if PDF could not be decrypted

        for page_num in range(len(reader.pages)):
            page = reader.pages[page_num]
            text += page.extract_text()

    if not text.strip():  # Debugging line to check if text is empty
        logger.warning(f"Could not decrypt PDF {file}")

    return text


def split_files_to_chunks(
    files: list[Union[tuple[str, str], str]],
    max_tokens: int = 4000,
    chunk_mode: str = "multi_lines",
    must_break_at_empty_line: bool = True,
    custom_text_split_function: Optional[Callable[[str], list[str]]] = None,
) -> tuple[list[str], list[dict[str, Any]]]:
    """Split a list of files into chunks of max_tokens."""
    chunks = []
    sources = []

    for file in files:
        if isinstance(file, tuple):
            url = file[1]
            file = file[0]
        else:
            url = None
        _, file_extension = os.path.splitext(file)
        file_extension = file_extension.lower()

        if HAS_UNSTRUCTURED and file_extension[1:] in UNSTRUCTURED_FORMATS:
            text = partition(file)
            text = "\n".join([t.text for t in text]) if len(text) > 0 else ""
        elif file_extension == ".pdf":
            text = extract_text_from_pdf(file)
        else:  # For non-PDF text-based files
            with open(file, encoding="utf-8", errors="ignore") as f:
                text = f.read()

        if not text.strip():  # Debugging line to check if text is empty after reading
            logger.warning(f"No text available in file: {file}")
            continue  # Skip to the next file if no text is available

        if custom_text_split_function is not None:
            tmp_chunks = custom_text_split_function(text)
        else:
            tmp_chunks = split_text_to_chunks(text, max_tokens, chunk_mode, must_break_at_empty_line)
        chunks += tmp_chunks
        sources += [{"source": url if url else file}] * len(tmp_chunks)

    return chunks, sources


def get_files_from_dir(
    dir_path: Union[str, list[str]], types: list[str] = TEXT_FORMATS, recursive: bool = True
) -> list[Any]:
    """Return a list of all the files in a given directory, a url, a file path or a list of them."""
    if len(types) == 0:
        raise ValueError("types cannot be empty.")
    types = [t[1:].lower() if t.startswith(".") else t.lower() for t in set(types)]
    types += [t.upper() for t in types]

    files = []
    # If the path is a list of files or urls, process and return them
    if isinstance(dir_path, list):
        for item in dir_path:
            if os.path.isfile(item):
                files.append(item)
            elif is_url(item):
                filepath = get_file_from_url(item)
                if filepath:
                    files.append(filepath)
            elif os.path.exists(item):
                try:
                    files.extend(get_files_from_dir(item, types, recursive))
                except ValueError:
                    logger.warning(f"Directory {item} does not exist. Skipping.")
            else:
                logger.warning(f"File {item} does not exist. Skipping.")
        return files

    # If the path is a file, return it
    if os.path.isfile(dir_path):
        return [dir_path]

    # If the path is a url, download it and return the downloaded file
    if is_url(dir_path):
        filepath = get_file_from_url(dir_path)
        if filepath:
            return [filepath]
        else:
            return []

    if os.path.exists(dir_path):
        for type in types:
            if recursive:
                files += glob.glob(os.path.join(dir_path, f"**/*.{type}"), recursive=True)
            else:
                files += glob.glob(os.path.join(dir_path, f"*.{type}"), recursive=False)
    else:
        logger.error(f"Directory {dir_path} does not exist.")
        raise ValueError(f"Directory {dir_path} does not exist.")
    return files


@require_optional_import(["markdownify", "bs4"], "retrievechat")
def parse_html_to_markdown(html: str, url: str = None) -> str:
    """Parse HTML to markdown."""
    soup = BeautifulSoup(html, "html.parser")
    title = soup.title.string
    # Remove javascript and style blocks
    for script in soup(["script", "style"]):
        script.extract()

    # Convert to markdown -- Wikipedia gets special attention to get a clean version of the page
    if isinstance(url, str) and url.startswith("https://en.wikipedia.org/"):
        body_elm = soup.find("div", {"id": "mw-content-text"})
        title_elm = soup.find("span", {"class": "mw-page-title-main"})

        if body_elm:
            # What's the title
            main_title = soup.title.string
            if title_elm and len(title_elm) > 0:
                main_title = title_elm.string
            webpage_text = "# " + main_title + "\n\n" + markdownify.MarkdownConverter().convert_soup(body_elm)
        else:
            webpage_text = markdownify.MarkdownConverter().convert_soup(soup)
    else:
        webpage_text = markdownify.MarkdownConverter().convert_soup(soup)

    # Convert newlines
    webpage_text = re.sub(r"\r\n", "\n", webpage_text)
    webpage_text = re.sub(r"\n{2,}", "\n\n", webpage_text).strip()
    webpage_text = "# " + title + "\n\n" + webpage_text
    return webpage_text


def _generate_file_name_from_url(url: str, max_length=255) -> str:
    url_bytes = url.encode("utf-8")
    hash = hashlib.blake2b(url_bytes).hexdigest()
    parsed_url = urlparse(url)
    file_name = os.path.basename(url)
    file_name = (
        f"{parsed_url.netloc}_{file_name}_{hash[: min(8, max_length - len(parsed_url.netloc) - len(file_name) - 1)]}"
    )
    return file_name


def get_file_from_url(url: str, save_path: str = None) -> tuple[str, str]:
    """Download a file from a URL."""
    if save_path is None:
        save_path = "tmp/chromadb"
        os.makedirs(save_path, exist_ok=True)
    if os.path.isdir(save_path):
        filename = _generate_file_name_from_url(url)
        save_path = os.path.join(save_path, filename)
    else:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

    custom_headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36",
    }
    try:
        response = requests.get(url, stream=True, headers=custom_headers, timeout=30)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        logger.warning(f"Failed to download {url}, {e}")
        return None

    content_type = response.headers.get("content-type", "")
    if "text/html" in content_type:
        # Get the content of the response
        html = ""
        for chunk in response.iter_content(chunk_size=8192, decode_unicode=True):
            html += chunk
        text = parse_html_to_markdown(html, url)
        with open(save_path, "w", encoding="utf-8") as f:
            f.write(text)
    else:
        with open(save_path, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
    return save_path, url


def is_url(string: str):
    """Return True if the string is a valid URL."""
    try:
        result = urlparse(string)
        return all([result.scheme, result.netloc])
    except ValueError:
        return False


@require_optional_import("chromadb", "retrievechat")
def create_vector_db_from_dir(
    dir_path: Union[str, list[str]],
    max_tokens: int = 4000,
    client: "API" = None,
    db_path: str = "tmp/chromadb.db",
    collection_name: str = "all-my-documents",
    get_or_create: bool = False,
    chunk_mode: str = "multi_lines",
    must_break_at_empty_line: bool = True,
    embedding_model: str = "all-MiniLM-L6-v2",
    embedding_function: Callable = None,
    custom_text_split_function: Callable = None,
    custom_text_types: list[str] = TEXT_FORMATS,
    recursive: bool = True,
    extra_docs: bool = False,
) -> "API":
    """Create a vector db from all the files in a given directory, the directory can also be a single file or a url to
        a single file. We support chromadb compatible APIs to create the vector db, this function is not required if
        you prepared your own vector db.

    Args:
        dir_path (Union[str, List[str]]): the path to the directory, file, url or a list of them.
        max_tokens (Optional, int): the maximum number of tokens per chunk. Default is 4000.
        client (Optional, API): the chromadb client. Default is None.
        db_path (Optional, str): the path to the chromadb. Default is "tmp/chromadb.db". The default was `/tmp/chromadb.db` for version `<=0.2.24`.
        collection_name (Optional, str): the name of the collection. Default is "all-my-documents".
        get_or_create (Optional, bool): Whether to get or create the collection. Default is False. If True, the collection
            will be returned if it already exists. Will raise ValueError if the collection already exists and get_or_create is False.
        chunk_mode (Optional, str): the chunk mode. Default is "multi_lines".
        must_break_at_empty_line (Optional, bool): Whether to break at empty line. Default is True.
        embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if
            embedding_function is not None.
        embedding_function (Optional, Callable): the embedding function to use. Default is None, SentenceTransformer with
            the given `embedding_model` will be used. If you want to use OpenAI, Cohere, HuggingFace or other embedding
            functions, you can pass it here, follow the examples in `https://docs.trychroma.com/embeddings`.
        custom_text_split_function (Optional, Callable): a custom function to split a string into a list of strings.
            Default is None, will use the default function in `autogen.retrieve_utils.split_text_to_chunks`.
        custom_text_types (Optional, List[str]): a list of file types to be processed. Default is TEXT_FORMATS.
        recursive (Optional, bool): whether to search documents recursively in the dir_path. Default is True.
        extra_docs (Optional, bool): whether to add more documents in the collection. Default is False

    Returns:
    The chromadb client.
    """
    if client is None:
        client = chromadb.PersistentClient(path=db_path)
    try:
        embedding_function = (
            ef.SentenceTransformerEmbeddingFunction(embedding_model)
            if embedding_function is None
            else embedding_function
        )
        collection = client.create_collection(
            collection_name,
            get_or_create=get_or_create,
            embedding_function=embedding_function,
            # https://github.com/nmslib/hnswlib#supported-distances
            # https://github.com/chroma-core/chroma/blob/566bc80f6c8ee29f7d99b6322654f32183c368c4/chromadb/segment/impl/vector/local_hnsw.py#L184
            # https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md
            metadata={"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32},  # ip, l2, cosine
        )

        length = 0
        if extra_docs:
            length = len(collection.get()["ids"])

        if custom_text_split_function is not None:
            chunks, sources = split_files_to_chunks(
                get_files_from_dir(dir_path, custom_text_types, recursive),
                custom_text_split_function=custom_text_split_function,
            )
        else:
            chunks, sources = split_files_to_chunks(
                get_files_from_dir(dir_path, custom_text_types, recursive),
                max_tokens,
                chunk_mode,
                must_break_at_empty_line,
            )
        logger.info(f"Found {len(chunks)} chunks.")
        # Upsert in batch of 40000 or less if the total number of chunks is less than 40000
        for i in range(0, len(chunks), min(40000, len(chunks))):
            end_idx = i + min(40000, len(chunks) - i)
            collection.upsert(
                documents=chunks[i:end_idx],
                ids=[f"doc_{j + length}" for j in range(i, end_idx)],  # unique for each doc
                metadatas=sources[i:end_idx],
            )
    except ValueError as e:
        logger.warning(f"{e}")
    return client


@require_optional_import("chromadb", "retrievechat")
def query_vector_db(
    query_texts: list[str],
    n_results: int = 10,
    client: "API" = None,
    db_path: str = "tmp/chromadb.db",
    collection_name: str = "all-my-documents",
    search_string: str = "",
    embedding_model: str = "all-MiniLM-L6-v2",
    embedding_function: Callable = None,
) -> "QueryResult":
    """Query a vector db. We support chromadb compatible APIs, it's not required if you prepared your own vector db
        and query function.

    Args:
        query_texts (List[str]): the list of strings which will be used to query the vector db.
        n_results (Optional, int): the number of results to return. Default is 10.
        client (Optional, API): the chromadb compatible client. Default is None, a chromadb client will be used.
        db_path (Optional, str): the path to the vector db. Default is "tmp/chromadb.db". The default was `/tmp/chromadb.db` for version `<=0.2.24`.
        collection_name (Optional, str): the name of the collection. Default is "all-my-documents".
        search_string (Optional, str): the search string. Only docs that contain an exact match of this string will be retrieved. Default is "".
        embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if
            embedding_function is not None.
        embedding_function (Optional, Callable): the embedding function to use. Default is None, SentenceTransformer with
            the given `embedding_model` will be used. If you want to use OpenAI, Cohere, HuggingFace or other embedding
            functions, you can pass it here, follow the examples in `https://docs.trychroma.com/embeddings`.

    Returns:
        The query result. The format is:

    ```python
    class QueryResult(TypedDict):
        ids: List[IDs]
        embeddings: Optional[List[List[Embedding]]]
        documents: Optional[List[List[Document]]]
        metadatas: Optional[List[List[Metadata]]]
        distances: Optional[List[List[float]]]
    ```
    """
    if client is None:
        client = chromadb.PersistentClient(path=db_path)
    # the collection's embedding function is always the default one, but we want to use the one we used to create the
    # collection. So we compute the embeddings ourselves and pass it to the query function.
    collection = client.get_collection(collection_name)
    embedding_function = (
        ef.SentenceTransformerEmbeddingFunction(embedding_model) if embedding_function is None else embedding_function
    )
    query_embeddings = embedding_function(query_texts)
    # Query/search n most similar results. You can also .get by id
    results = collection.query(
        query_embeddings=query_embeddings,
        n_results=n_results,
        where_document={"$contains": search_string} if search_string else None,  # optional filter
    )
    return results