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 time
import typing
import wikipedia as wiki
from .constants import ENV, LLM_API_KEY, LLM_BASE_URL, MEILI_URL, MEILI_MASTER_KEY, INDEX

WIKI_ATTRIBUTES = 'content images title sections html links summary pageid url revision_id parent_id categories'.split()
WIKI_SECTION_ATTRIBUTES = 'title url revision_id parent_id categories'.split()
PREFIX_STOPWORDS = 'who what when where why how it is are do does am I to the a an'.split()
LLM_NAME = 'ggml-org/gemma-3-1b-it-gguf'

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, **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,
    })
    response = client.chat.completions.create(
        model=llm_name,
        messages=messages,
        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
):
    """ 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


def main():
    import sys
    print(ask(' '.join(sys.argv[1:])))
