from typing import List, Dict, Optional, Union, Any
from alana.color import red, yellow
import re
import os
from anthropic import Anthropic
from anthropic.types import MessageParam

def get_xml_pattern(tag: str):
    if tag.count('<') > 0 or tag.count('>') > 0:
        raise ValueError("No '>' or '<' allowed in get_xml tag name!")
    return rf"<{tag}>(.*?)</{tag}>"

def get_xml(tag: str, content: str) -> List[str]:
    if tag.count('<') > 0 or tag.count('>') > 0:
        raise ValueError("No '>' or '<' allowed in get_xml tag name!")
    pattern = get_xml_pattern(tag)
    matches = re.findall(pattern, content, re.DOTALL)
    return matches

DEFAULT_MODEL = "opus"
MODELS = {
    'opus' : 'claude-3-opus-20240229',
    'claude-3-opus-20240229' : 'claude-3-opus-20240229',
    'sonnet' : 'claude-3-sonnet-20240229',
    'claude-3-sonnet-20240229' : 'claude-3-sonnet-20240229',
    'haiku' : 'claude-3-haiku-20240307',
    'claude-3-haiku-0307' : 'claude-3-haiku-20240307',
    'claude-2.1' : 'claude-2.1',
    'claude-2.0' : 'claude-2.0',
    'claude-instant-1.2' : 'claude-instant-1.2'
}


def gen(user: Optional[str] = None, system: str = "", messages: Optional[List[MessageParam]] = None, model: str = DEFAULT_MODEL, api_key: Optional[str] = None, max_tokens = 1024, temperature=0.3, quiet=False, **kwargs) -> str:
    if user is None and messages is None:
        raise ValueError("No prompt provided! `user` and `messages` are both None.")
    elif messages is None:
        messages=[
            MessageParam(
                role="user",
                content=user, # type: ignore
            ),
        ]

    backend = MODELS[DEFAULT_MODEL]
    if model in MODELS:
        backend = MODELS[model]
    else:
        red(f"gen() -- Caution! model string not recognized; reverting to {DEFAULT_MODEL=}.") # TODO: C'mon we can do better error logging than this

    if api_key is None:
        api_key = os.environ.get("ANTHROPIC_API_KEY")
    client = Anthropic(
        api_key=api_key,
    )

    message = client.messages.create(
        max_tokens=max_tokens,
        messages=messages,
        system=system,
        temperature=temperature,
        model=backend,
        **kwargs
    )
    if not quiet:
        yellow(message)
    return message.content[0].text

def gen_examples_list(instruction: str, n_examples: int = 5, model: str = "sonnet", api_key: Optional[str] = None, max_tokens: int = 1024, temperature=0.3, **kwargs) -> List[str]:
    system = """You are a prompt engineering assistant tasked with generating few-shot examples given a task.

    Your user would like to accomplish a specific task. His/her description of the task will be enclosed in <description></description> XML tags.
    
    You are to generate {n_examples} examples of this task. EACH example must be enclosed in <example></example> XML tags. Each example should make the input and output clear.

    Make sure your examples are clear, high-quality, consistent, and cover a range of use-cases.
    """.format(n_examples=n_examples)

    if n_examples < 1:
        red("Too few examples provided! Trying anyway...")

    user = """The user's task is as follows:
    <description>{instruction}</description>

    Before generating your examples, feel free to think out loud using <thinking></thinking> XML tags.
    """.format(instruction=instruction)
    model_output = gen(user=user, system=system, model=model, api_key=api_key, max_tokens=max_tokens, temperature=temperature, **kwargs)
    return get_xml('example', model_output)

def gen_examples(instruction: str, n_examples: int = 5, model: str = DEFAULT_MODEL, api_key: Optional[str] = None, max_tokens: int = 1024, temperature=0.3, **kwargs) -> str:
    examples = gen_examples_list(instruction=instruction, n_examples=n_examples, model=model, api_key=api_key, max_tokens=max_tokens, temperature=temperature, **kwargs)
    formatted_examples = "\n<examples>\n<example>" + '</example>\n<example>'.join(examples) + "</example>\n</examples>"
    return formatted_examples

def gen_prompt(instruction: str, model: str = DEFAULT_MODEL, api_key: Optional[str] = None, max_tokens: int = 1024, temperature=0.3, **kwargs) -> Dict[str, Union[str, List]]:
    meta_system_prompt = """You are a prompt engineering assistant. Your task is to effectively prompt a language model to complete a given task.

    The task description will be enclosed in <description></description> XML tags. Your user may also provide important context in the description.

    You will generate both a system prompt and a user prompt. Depending on the task, you may decide to leave the system prompt empty.

    When writing the system prompt, consider the following components:
    1. Precise language. Make sure to clearly and precisely describe the task, and eliminate any room for misunderstanding.
    2. Context. Provide any important context for the task.
    3. Role (optional). Consider telling the model to inhabit a role (e.g. an expert programmer) to improve its response.

    When writing the user prompt, consider the following components:
    1. Input data. If the user intends to provide any data to the model, use a placeholder {like so}. Surround the input data with <user_input></user_input> XML tags. The user will provide the data using the Python string.format function.
    2. Clear instructions. Consider writing step-by-step instructions for the model to follow to complete the task.
    3. Output formatting. Specify the final output format that the model should conform to.
    4. Few-shot examples (optional). You might include some examples of the intended behavior. Enclose each example in <example></example> XML tags.
    5. "Prompting tricks." Consider asking the model to think out loud. Consider writing particularly important phrases in ALL CAPS.

    Be careful:
    1. You MUST enclose your system prompt in <system_prompt></system_prompt> XML tags.
    2. You MUST enclose your user prompt in <user_prompt></user_prompt> XML tags.

    Before producing your prompt, feel free to think out loud using <reasoning></reasoning> XML tags. Enclose your thinking in <reasoning></reasoning> XML tags.
    """
    meta_prompt = """Here is the task description:

    <description>
    {instruction}
    </description>

    Now:
    1. Feel free to think out loud in <reasoning></reasoning> XML tags.
    2. Enclose your system prompt in <system_prompt></system_prompt> XML tags.
    3. Enclose your user prompt in <user_prompt></user_prompt> XML tags.
    """.format(instruction=instruction)

    full_output = gen(user=meta_prompt, system=meta_system_prompt, model=model, api_key=api_key, max_tokens=max_tokens, temperature=temperature, **kwargs)
    system_prompt = get_xml("system_prompt", full_output)
    if len(system_prompt) >= 1:
        system_prompt = system_prompt[0]
    user_prompt = get_xml("user_prompt", full_output)
    if len(user_prompt) == 1:
        user_prompt = user_prompt[0]
    return {"system": system_prompt, "user": user_prompt, "full": full_output}

def pretty_print(var: Any, loud: bool = True) -> str:
    raise NotImplementedError

# Aliases!
def grab(tag: str, content: str) -> List[str]:
    """Alias for get_xml"""
    return get_xml(tag, content)

def xml(tag: str, content: str) -> List[str]:
    """Alias for get_xml"""
    return get_xml(tag, content)

def n_shot_list(instruction: str, n_examples: int = 5, model: str = DEFAULT_MODEL, api_key: Optional[str] = None, max_tokens: int = 1024, temperature=0.3, **kwargs) -> List[str]:
    """Alias for gen_examples_list"""
    return gen_examples_list(instruction=instruction, n_examples=n_examples, model=model, api_key=api_key, max_tokens=max_tokens, temperature=temperature, **kwargs)

def n_shot(instruction: str, n_examples: int = 5, model: str = DEFAULT_MODEL, api_key: Optional[str] = None, max_tokens: int = 1024, temperature=0.3, **kwargs) -> str:
    """Alias for gen_examples"""
    return gen_examples(instruction=instruction, n_examples=n_examples, model=model, api_key=api_key, max_tokens=max_tokens, temperature=temperature, **kwargs)

def few_shot_list(instruction: str, n_examples: int = 5, model: str = DEFAULT_MODEL, api_key: Optional[str] = None, max_tokens: int = 1024, temperature=0.3, **kwargs) -> List[str]:
    """Alias for gen_examples_list"""
    return gen_examples_list(instruction=instruction, n_examples=n_examples, model=model, api_key=api_key, max_tokens=max_tokens, temperature=temperature, **kwargs)

def few_shot(instruction: str, n_examples: int = 5, model: str = DEFAULT_MODEL, api_key: Optional[str] = None, max_tokens: int = 1024, temperature=0.3, **kwargs) -> str:
    """Alias for gen_examples"""
    return gen_examples(instruction=instruction, n_examples=n_examples, model=model, api_key=api_key, max_tokens=max_tokens, temperature=temperature, **kwargs)
