import dotenv
import json
import logging
from meilisearch_python_sdk import Client, Index, errors
from markdown_chunker import MarkdownChunkingStrategy
import os
from openai import OpenAI
from pathlib import Path
import re
import string
import sys
import time
import typing
import wikipedia as wiki
from .constants import ENV, LLM_API_KEY, LLM_BASE_URL, MEILI_URL, MEILI_MASTER_KEY, INDEX

from .constants import WIKI_ATTRIBUTES, WIKI_SECTION_ATTRIBUTES, PREFIX_STOPWORDS, LLM_NAME
SYSTEM_PROMPT = "You are a factual question-answering chatbot within a RAG."

log = logging.getLogger()


def make_pk(text: str):
    """ Only letters, numbers, _, and - are allowed in docuent ID's (PK's) """
    text = text.replace(' ', '_')
    text = text.replace('(', '-_')
    text = text.replace(')', '_-')
    text = ''.join(c for c in text if c in string.ascii_letters + string.digits + '-_')
    return text


def clean_text(text: str):
    """ Remove most punctuation for use simple text classification and keyword detection

    >>> clean_text('2,000')
    '2000'
    >>> clean_text('hello-{WORLD}!')
    'hello-_WORLD_'
    """
    text = re.sub(r'[[\](){}]+', '_', text)
    text = re.sub(r'[^-\w\s]+', ' ', text)
    return text.strip()


def make_dict(page, attributes=WIKI_ATTRIBUTES):
    if isinstance(page, str):
        page = wiki.page(page)
    page_dict = {k: getattr(page, k) for k in attributes if hasattr(page, k)}
    page_dict.update({k: v() for (k, v) in page_dict.items() if callable(v)})
    return page_dict


def make_name(s: str):
    r""" Slugify a short str: delete whitespace, replace non-alphanum (except _) with "-" """
    s = re.sub(r'\s+', '', s)
    s = s.replace('.', '_')
    s = s.replace(':', '-')
    s = re.sub(r'[^-_0-9a-zA-Z]+', '-', s)
    return s


def make_index(name_or_index=None, url=MEILI_URL, api_key=MEILI_MASTER_KEY):
    global INDEX
    name_or_index = name_or_index or INDEX
    if isinstance(name_or_index, Index):
        INDEX = name_or_index
        return INDEX
    log.info(f'Creating meilisearch.client.Client("{url}")')
    c = Client(url=url, api_key=api_key)
    if isinstance(name_or_index, (str, int, float)):
        name = make_name(str(name_or_index))
        try:
            INDEX = c.get_index(name)
            return INDEX
        except errors.MeilisearchError as e:
            log.warning(f'Unable to meili.get_index("{name}"): {e}')
    c.create_index(uid=name, primary_key='pk', wait=True, hits_type='Any')
    for i in range(5):
        time.sleep(1.5**i)
        INDEX = c.get_index(name)
        if INDEX:
            return INDEX


def make_md_chunks(text):
    strategy = MarkdownChunkingStrategy(add_metadata=True)
    return strategy.chunk_markdown(text)


def make_wiki_chunks(page, attributes=WIKI_SECTION_ATTRIBUTES):
    if not isinstance(page, dict):
        page = make_dict(page)
    chunks = []
    if not 'sections' in page:
        return page

    sections = page['sections'].copy()
    del page['sections']
    chunks = []
    for i, (k, v) in enumerate(sections.items()):
        chnk = dict(v).copy()
        pk = make_pk(str(page['title']) + '_--_' + str(k))
        log.debug(f'composing chunk for pk: {pk}')
        chnk.update({sect_k: v for sect_k, v in page.items() if sect_k in attributes})
        chnk['heading'] = str(k)
        chnk['pk'] = pk
        chunks.append(chnk)
    log.debug(f'len(make_wiki_chunks(page={str(page)[:50]}...)): {len(chunks)}')
    return chunks


def make_wiki_title(title):
    return title.lower().replace('-', ' ').replace('_', ' ')


def add_chunks(chunks: [str | wiki.WikipediaPage | list[dict]]):
    """ Chunk a Wikipedia article and add it to the meilisearch index

    >>> add_chunks('Amylose')
    <class Index...>
    >>> add_chunks(wiki.page('Glucose'))
    <class Index...>
    >>> add_chunks(make_wiki_chunks(wiki.page('Insulin')))
    <class Index...>
    """
    INDEX = make_index()
    if isinstance(chunks, (str, wiki.WikipediaPage)):
        log.debug(f'Creating chunks from wikipedia.Page: {chunks}')
        chunks = make_wiki_chunks(chunks)
    INDEX.add_documents(chunks, primary_key='pk')
    log.debug(f'ragkit.util.INDEX: {INDEX}')
    return INDEX


def add_wiki_pages(
        pages: [str | list | wiki.WikipediaPage],
        depth: int = 0):
    """ Chunkify a list of wikipedia pages and add them to the default meilisearch index

    TODO: May be redundant with `add_chunks()` etc

    >>> add_wiki_pages('Virus', depth=1)
    <class Index...
    >>> add_wiki_pages(['Gluose'])
    [<class Index...]
    >>> add_wiki_pages(['Gluose', 'Sucrose', 'Insulin'])
    [<class Index...]
    """
    if isinstance(pages, str):
        pages = make_wiki_title(pages)
        try:
            page = wiki.get_page(pages)
        except Exception as e:
            log.error(f'Could not find "{pages}"" in add_wiki_pages(pages="{pages}", depth={depth})')
            return None
        if depth:
            return add_wiki_pages([pages] + list(page.links), depth=depth - 1)
        return add_wiki_pages(page, depth=depth)

    if isinstance(pages, wiki.WikipediaPage):
        log.debug(f'Creating dict and chunks from wikipedia.Page: {pages}')
        page_dict = make_dict(pages)
        chunks = make_wiki_chunks(page_dict)
        with open('cache-wikipedia.jsonlines', 'a') as fapp:
            fapp.write(json.dumps(page_dict) + '\n')
            fapp.flush()
        return add_chunks(chunks)

    if isinstance(pages, list) and len(pages):
        if isinstance(pages[0], str):
            return [add_wiki_pages(title, depth=depth) for title in pages]
        return add_chunks(pages)

    log.error(f'Bad argument types in add_wiki_pages(pages="{pages}", depth={depth})')
    return None


def updated(d: dict, update: dict):
    """ DEPRECATED: inline this pattern """
    return d.update(u) or d


def search(query: str, limit: int = 3, facets=None, score=True, hits_per_page=20, **opt_params) -> list[dict]:
    """ Retrieve document dicts from default meilisearch INDEX and return a list of hits"""
    opt_params.update(dict(
        limit=limit,
        facets=facets,
        show_ranking_score=score,
        hits_per_page=hits_per_page,
    )
    )
    INDEX = make_index()
    res = INDEX.search(query, **opt_params)
    for hit in res.hits:
        hit.update(dict(num_hits=res.total_hits or res.estimated_total_hits, offset=res.offset))
    return res.hits


def classify_text(query: str):
    """ Classify user search query (prompt) as prompt, search, advanced, sql, OR empty

    An 'advanced' query is a natural language search request that contains OR or AND logic
    Class is typically used to route the user query directly to a search engine, SQL parser, or LLM/RAG

    >>> classify_text('What is sucrose?')
    'prompt'
    >>> classify_text('sucrose glucose insulin')
    'search'
    >>> classify_text('SELECT name FROM Food where glycemic_index > 30 LIMIT 3')
    'sql'
    >>> classify_text('How big is a typical virus?')
    'prompt'
    """
    query = query.strip()
    prefix = query[:4].lower()
    if query.endswith('?') or prefix == '@rag':
        return 'rag'
    if prefix == '@llm':
        return 'llm'
    if query.startswith('@sql'):
        return 'sql'
    query = clean_text(query)
    if not query or not query.strip():
        return 'empty'
    if len(query) > 10:
        if ' AND ' in query or ' OR ' in query:
            return 'advanced'  # advanced search
        words = query.lower().split()
        if len(words) > 5:
            return 'rag'
        if len(words) > 3 and len(words) < 32 and words[0] in ('from', 'sql', 'psql') and ('select' in words or 'where' in words):
            return 'sql'
    return 'search'


def make_search_query(text):
    words = text.lower().split()
    for psw in PREFIX_STOPWORDS:
        if words[0] == psw:
            text = text[len(psw):].lstrip()
            words = words[1:]
    return text


def make_rag_prompt(query: str, results: list, limit=1, template_name=None, **kwargs):
    if not template_name:
        prompt = f'''
        Only use only the following pieces of <CONTEXT> to answer the <QUESTION>.
        Don't make up any new information.

        '''
        prompt += '<CONTEXT>\n'
        prompt += json.dumps([{ky: results[i][ky] for ky in 'title heading text'.split()}
                              for i in range(limit)], indent=4).strip()
        prompt += '\n</CONTEXT>\n\n'
        prompt += '<QUESTION>\n' + query.strip() + '\n</QUESTION>\n\n'

    # print('## PROMPT:')
    # print(prompt)

    return prompt


def complete_with_llm(prompt, system_prompt=SYSTEM_PROMPT, max_tokens=256, base_url=LLM_BASE_URL,
                      api_key="sk-xxx", llm_name=LLM_NAME, answer_only=True, temperature=0.3, cache_prompt=False, **kwargs):
    client = OpenAI(base_url=base_url, api_key=api_key)
    messages = []
    if system_prompt.strip():
        messages.append({
            "role": "system",
            "content": system_prompt,
        })
    messages.append({
        "role": "user",
        "content": prompt,
    })
    extra_body = kwargs.pop('extra_body', {"cache_prompt": cache_prompt})
    extra_body.update(kwargs)
    response = client.chat.completions.create(
        model=llm_name,
        messages=messages,
        extra_body=extra_body,
        max_tokens=max_tokens,
        temperature=temperature,
        logprobs=False,
        **kwargs,
    )
    # print('## LLM RESPONSE OBJECT:')
    # print(response)

    answer = response.choices[0].message.content
    if answer_only:
        return answer
    response = response.dict()
    response.update(dict(prompt=prompt, answer=answer))
    return response


def ask(
        query: str,
        limit: int = 3,
        max_tokens: int = 128,
        top_p: float = 0.2,
        temperature: float = 0.2,
        base_url: str = LLM_BASE_URL,
        llm_name: str = LLM_NAME,
        answer_only: bool = True,
        automatic_routing: bool = True,
        **kwargs,
):
    """ Prompt an LLM web API with search results to answer a question 

    curl --request POST \
        --url http://localhost:8080/completion \
        --header "Content-Type: application/json" \
        --data '{"prompt": "Can a RAG answer questions?", "n_predict": 128}'
    """
    search_query = make_search_query(query)
    results = search(search_query, limit=limit)
    if automatic_routing and classify_text(query) == 'search':
        return results
    rag_prompt = make_rag_prompt(query, results=results, limit=1)
    answer = complete_with_llm(rag_prompt, base_url=base_url, llm_name=llm_name,
                               answer_only=answer_only, max_tokens=max_tokens,
                               top_p=top_p, temperature=temperature)
    return answer
