import logging
import os
import subprocess as sp
from typing import Iterator, List, Dict, Tuple, Generator
from urllib.parse import urlparse

import aiohttp
import numpy as np
import yaml

from llama_cpp import Llama, CreateCompletionResponse, CreateCompletionStreamResponse
from logdetective.constants import SNIPPET_DELIMITER
from logdetective.models import PromptConfig, SkipSnippets
from logdetective.remote_log import RemoteLog

LOG = logging.getLogger("logdetective")


def new_message(text: str) -> bool:
    """Set of heuristics for determining whether or not
    does the current chunk of log text continue on next line.

    Following rules are checked, in order:
    * is the first character is whitespace
    * is the first character backslash '|'
    """
    conditionals = [
        lambda string: string[0].isspace(),
        lambda string: string[0] == "|",
    ]

    for c in conditionals:
        y = c(text)
        if y:
            return False

    return True


def get_chunks(
    text: str, max_chunk_len: int = 2000
) -> Generator[Tuple[int, str], None, None]:
    """Split log into chunks according to heuristic
    based on whitespace and backslash presence.
    """
    lines = text.splitlines()

    # Chunk we will be yielding
    chunk = ""
    # Number of line where the message started
    original_line = 0
    for i, line in enumerate(lines):
        if len(line) == 0:
            continue
        if new_message(line):
            # Yield chunk if we have it
            if len(chunk) > 0:
                yield (original_line, chunk)
            original_line = i
            chunk = line
        else:
            chunk += "\n" + line
        if len(chunk) > max_chunk_len:
            # If the chunk is too long, keep splitting into smaller chunks
            # until we reach manageable size
            while len(chunk) > max_chunk_len:
                remainder = chunk[max_chunk_len:]
                chunk = chunk[:max_chunk_len]
                yield (original_line, chunk)
                chunk = remainder

    # if we still have some text left over
    yield (original_line, chunk)


def initialize_model(
    model_pth: str, filename_suffix: str = ".gguf", verbose: bool = False
) -> Llama:
    """Initialize Llama class for inference.
    Args:
        model_pth (str): path to gguf model file or Hugging Face name
        filename_suffix (str): suffix of the model file name to be pulled from Hugging Face
        verbose (bool): level of verbosity for llamacpp
    """

    LOG.info("Loading model from %s", model_pth)

    if os.path.isfile(model_pth):
        model = Llama(
            model_path=model_pth,
            n_ctx=0,  # Maximum context for the model
            verbose=verbose,
            logits_all=True,
        )
    else:
        model = Llama.from_pretrained(
            model_pth,
            f"*{filename_suffix}",
            n_ctx=0,  # Maximum context for the model
            verbose=verbose,
            logits_all=True,
        )

    return model


def compute_certainty(probs: List[Dict]) -> float:
    """Compute certainty of repsponse based on average logit probability.
    Log probability is log(p), isn't really readable for most people, especially in compound.
    In this case it's just a matter of applying inverse operation exp.
    Of course that leaves you with a value in range <0, 1> so it needs to be multiplied by 100.
    Simply put, this is the most straightforward way to get the numbers out.

    This function is used in the server codebase.
    """

    top_logprobs = [np.exp(e["logprob"]) * 100 for e in probs]

    certainty = np.median(top_logprobs, axis=0)
    if np.isnan(certainty):
        raise ValueError("NaN certainty of answer")
    return certainty


def process_log(
    log: str, model: Llama, stream: bool, prompt_template: str, temperature: float
) -> CreateCompletionResponse | Iterator[CreateCompletionStreamResponse]:
    """Processes a given log using the provided language model and returns its summary.

    Args:
        log (str): The input log to be processed.
        model (Llama): The language model used for processing the log.
        stream (bool): Return output as Iterator.
        prompt_template (str): Which prompt template to use.
        temperature (float): Temperature parameter for model runtime.
    Returns:
        str: The summary of the given log generated by the language model.
    """
    response = model(
        prompt=prompt_template.format(log),
        stream=stream,
        max_tokens=0,
        logprobs=1,
        temperature=temperature,
    )

    return response


async def retrieve_log_content(http: aiohttp.ClientSession, log_path: str) -> str:
    """Get content of the file on the log_path path.
    Path is assumed to be valid URL if it has a scheme.
    Otherwise it attempts to pull it from local filesystem."""
    parsed_url = urlparse(log_path)
    log = ""

    if not parsed_url.scheme:
        if not os.path.exists(log_path):
            raise ValueError(f"Local log {log_path} doesn't exist!")

        with open(log_path, "rt") as f:
            log = f.read()

    else:
        remote_log = RemoteLog(log_path, http)
        log = await remote_log.get_url_content()

    return log


def format_snippets(snippets: list[str] | list[Tuple[int, str]]) -> str:
    """Format snippets, giving them separator, id and finally
    concatenating them. If snippets have line number attached,
    include that in prompt.

    Line number must be first element in the tuple. Mixed format of snippets
    is permitted, but may have impact on inference.
    """
    summary = ""
    for i, s in enumerate(snippets):
        if isinstance(s, tuple):
            summary += f"""
            Snippet No. {i} at line #{s[0]}:

            {s[1]}
            {SNIPPET_DELIMITER}
            """
        else:
            summary += f"""
            Snippet No. {i}:

            {s}
            {SNIPPET_DELIMITER}
            """
    return summary


def load_prompts(path: str | None) -> PromptConfig:
    """Load prompts from given yaml file if there is one.
    Alternatively use defaults."""
    if path:
        try:
            with open(path, "r") as file:
                return PromptConfig(yaml.safe_load(file))
        except FileNotFoundError:
            print("Prompt configuration file not found, reverting to defaults.")
    return PromptConfig()


def prompt_to_messages(
    user_message: str,
    system_prompt: str | None = None,
    system_role: str = "developer",
    user_role: str = "user",
) -> List[Dict[str, str]]:
    """Turn prompt into list of message dictionaries.
    If `system_role` and `user_role` are the same, only a single message is created,
    as concatenation of `user_message` and `system_prompt`. This is useful for models which
    do not have separate system role, such as mistral.
    """

    if system_role == user_role:
        messages = [
            {"role": system_role, "content": f"{system_prompt}\n{user_message}"}
        ]
    else:
        messages = [
            {"role": system_role, "content": system_prompt},
            {
                "role": user_role,
                "content": user_message,
            },
        ]

    return messages


def filter_snippet_patterns(snippet: str, skip_snippets: SkipSnippets) -> bool:
    """Try to match snippet agains provided patterns to determine if we should
    filter it out or not."""
    for key, pattern in skip_snippets.snippet_patterns.items():
        if pattern.match(snippet):
            LOG.debug("Snippet `%s` has matched agains skip pattern %s", snippet, key)
            return True

    return False


def load_skip_snippet_patterns(path: str | None) -> SkipSnippets:
    """Load dictionary of snippet patterns we want to skip."""
    if path:
        try:
            with open(path, "r") as file:
                return SkipSnippets(yaml.safe_load(file))
        except OSError as e:
            LOG.error("Couldn't open file with snippet skip patterns `%s`", path)
            raise e

    return SkipSnippets({})


def check_csgrep() -> bool:
    """Verifies presence of csgrep in path"""
    try:
        result = sp.run(
            ["csgrep", "--version"],
            text=True,
            check=True,
            shell=False,
            capture_output=True,
            timeout=1.0,
        )
    except (FileNotFoundError, sp.TimeoutExpired, sp.CalledProcessError) as ex:
        LOG.error("Required binary `csgrep` was not found in path: %s", ex)
        return False
    if result.returncode == 0:
        return True
    LOG.error("Issue was encountered while calling `csgrep`: `%s`", result.stderr)

    return False


def mine_logs(log: str, extractors: list) -> List[Tuple[int, str]]:
    """Extract snippets from log text using extractors provided.
    Each extractor is applied in turn on original log.
    Depending on characteristics of extractors used, there may be
    an overlap in snippets extracted."""

    log_summary = []

    LOG.info("Getting summary")

    for extractor in extractors:
        log_summary.extend(extractor(log))

    ratio = len("\n".join([text for _, text in log_summary])) / len(log)
    LOG.debug("Log summary: \n %s", log_summary)
    LOG.info("Snippets: %s Compression ratio: %s", len(log_summary), ratio)

    return log_summary
