"""Retrieve documentation for a given query."""

from pathlib import Path
from typing import Any
from rich.console import Console
from tqdm import tqdm
import numpy as np
from manifest import Manifest
from langchain.text_splitter import MarkdownHeaderTextSplitter
from langchain.text_splitter import RecursiveCharacterTextSplitter

console = Console(soft_wrap=True)

try:
    EMBEDDING_MODEL = Manifest(
        client_name="openaiembedding",
    )
except Exception as e:
    console.print(e)
    console.print(
        "Failed to load embedding model. Likely OPENAI API key is not set. Please set to run document retrieval.",
        style="bold red",
    )


def load_documentation(path: Path) -> dict[str, str]:
    """Load documentation from path."""
    content = {}
    for file in path.glob("**/*.md"):
        with open(file, "r") as f:
            data = f.read()
            key = str(file).replace(str(path), "")
            content[key] = data
    return content


def split_documents(content: dict[str, str]) -> dict[str, Any]:
    """Split documents into chunks."""
    md_splitted_docs = []
    markdown_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=[("#", "Header 1"), ("##", "Header 2"), ("###", "Header 3")]
    )
    text_splitter = RecursiveCharacterTextSplitter(
        separators=["\n"], chunk_size=500, chunk_overlap=50, length_function=len
    )

    for file, raw_doc in content.items():
        splitted_text = markdown_splitter.split_text(raw_doc)
        for t in splitted_text:
            t.metadata["source"] = file
        md_splitted_docs.extend(splitted_text)

    docs = text_splitter.split_documents(md_splitted_docs)
    docs_as_dict = [doc.dict() for doc in docs]
    return docs_as_dict


def get_embeddings(text: str) -> np.ndarray:
    """Get embeddings."""
    return np.array(EMBEDDING_MODEL.run(text))


def embed_documents(
    chunked_docs: dict[str, Any], key: str = "page_content"
) -> tuple[dict[str, Any], np.ndarray]:
    """Embed documents."""
    all_embeddings = []
    for doc in tqdm(chunked_docs):
        emb = get_embeddings(doc[key])
        doc["embedding"] = emb
        all_embeddings.append(doc["embedding"])
    full_embedding_mat = np.vstack(all_embeddings)
    return chunked_docs, full_embedding_mat


def query_docs(
    query: str,
    docs: dict[str, Any],
    embedding_mat: np.ndarray,
    top_n: int = 10,
    key: str = "page_content",
) -> tuple[list[int], list[str]]:
    """Query documents."""
    query_embedding = get_embeddings(query)
    scores = embedding_mat.dot(query_embedding)
    sorted_indices = np.argsort(scores)[::-1]
    top_n_indices = sorted_indices[:top_n]
    top_n_indices_rev = top_n_indices[::-1]
    returned_docs = []
    for i in top_n_indices_rev:
        returned_docs.append(docs[i][key])
    return top_n_indices_rev.tolist(), returned_docs