import contextlib
import itertools
import logging
import time
from dataclasses import dataclass
from typing import Any, Generator, Optional

import tiktoken
from asgiref.sync import sync_to_async
from django.db import connections as db_connections
from django.http import HttpResponse, QueryDict
from django.template.loader import render_to_string

log = logging.getLogger(__name__)

# load gpt-3.5 encoding
encoding = tiktoken.encoding_for_model("gpt-3.5-turbo-0613")


def sse_message(data: Optional[str] = None, event: Optional[str] = None) -> bytes:
    """Format some data as a Server Sent Event (SSE) message.

    Useful inside async functions that stream data to the client using
    server-sent events. e.g. you can use `yield sse_message(data='Hello!',
    event='greeting')` to say hi.

    This escapes newlines, but DOES NOT apply HTML escaping, so be careful when
    using it to stream responses to HTMX."""
    message = ""
    if event is not None:
        # Escape newline characters in event.
        event = event.replace("\n", "\nevent: ")
        message += f"event: {event}\n"
    if data is not None:
        # Escape newline characters in data.
        data = data.replace("\n", "\ndata: ")
        message += f"data: {data}\n"
    message += "\n"  # Messages are separated by an extra newline.
    return message.encode("utf-8")


def render_html_response(template_name: str, **context: Any) -> HttpResponse:
    """Render a HTML template and return it as a response.

    Response is always HTML and always uses the 200 error code. Provided kwargs
    are passed into the template rendering function as context.

    Useful for sending stuff to HTMX."""
    html_content = render_to_string(template_name, context)
    return HttpResponse(html_content)


def is_null_or_white_space(candidate_str: Optional[str]) -> bool:
    return not candidate_str or candidate_str.isspace()


@dataclass
class TokenCheck:
    verification_passed: bool
    num_tokens: int


def get_number_tokens(message: str, limit: int, processing_size=500) -> TokenCheck:
    """Check the number of tokens while avoiding horrible quadratic behavior on
    inputs >processing_size tokens long.

    (This is still quadratic for <=processing_size tokens, but in practice that
    is pretty fast.)"""
    for end_index in range(0, len(message), processing_size):
        if len(encoding.encode(message[:end_index], disallowed_special=())) > limit:
            return TokenCheck(verification_passed=False, num_tokens=limit * 2)
    num_tokens = len(encoding.encode(message, disallowed_special=()))
    if num_tokens > limit:
        return TokenCheck(verification_passed=False, num_tokens=limit * 2)
    else:
        return TokenCheck(verification_passed=True, num_tokens=num_tokens)


def remove_gpt_token_runs(message: str) -> str:
    """Remove runs of GPT tokens from a message.

    This is a hacky workaround for the GPT 3.5 Turbo token repetition rejection
    bug where the OpenAI API typically rejects requests with runs of >=200 of
    the same token.

    Assumes that length has already been checked so we don't need to worry about
    quadratic behavior."""
    MAX_RUN_LEN = 199
    # tokens is a list[int]
    tokens = encoding.encode_ordinary(message)
    # create a new token list where no run is more than 199 tokens long
    new_tokens = []
    removed_run = False  # did we ever remove a run?
    run_len = 0  # length of current run
    for i in range(len(tokens)):
        if i > 0 and tokens[i] != tokens[i - 1]:
            # reset the run
            run_len = 0
        # add this token to the run
        run_len += 1
        if run_len <= MAX_RUN_LEN:
            # we're in a run that's short enough, so add the token
            new_tokens.append(tokens[i])
        else:
            # did we ever remove a run?
            removed_run = True

    if removed_run:
        # log if we ever had to remove anything (should be rare)
        log.info(
            "remove_gpt_token_runs had to remove a long run of tokens from a message"
        )

    return encoding.decode(new_tokens)


async def async_close_db_connections() -> None:
    """Close Django DB connections that can be closed.

    This closes db connections both in this thread and the background thread
    where sync_to_async runs. AFAICT Django should _only_ create connections in
    the background thread where sync_to_async runs, but I'm not 100% sure, so
    I'm doing this defensively.

    This is useful for long-lived async views that don't need to make a lot of
    database queries, since it prevents the connection objects from sticking
    around forever."""
    db_connections.close_all()
    await sync_to_async(db_connections.close_all)()


@contextlib.contextmanager
def log_timer(log_key: str) -> Generator[None, None, None]:
    log.info(f"Beginning timing of {log_key}")

    start = time.monotonic()
    yield  # let the thing run
    elapsed = time.monotonic() - start

    log.info({log_key: elapsed, "message": f"Ending timing of {log_key}"})


def short_format_dict(
    d: dict | QueryDict,
    *,
    key_length_limit: int = 20,
    value_length_limit: int = 50,
    n_items_limit=10,
) -> str:
    """Concisely string-format a dictionary.

    Args:
        d: The dictionary to string-format.
        key_length_limit: The maximum length of a key to include in the output.
        value_length_limit: The maximum length of a value to include in the output.
        n_items_limit: The maximum number of items to include in the output.

    Returns: A string representation of the dictionary."""

    def format_limited(s: str, limit: int) -> str:
        if len(s) > limit:
            return s[:limit] + "..."
        else:
            return s

    return (
        "{"
        + ", ".join(
            f"{format_limited(repr(k), key_length_limit)}: {format_limited(repr(v), value_length_limit)}"
            for k, v in itertools.islice(d.items(), n_items_limit)
        )
        + "}"
    )
