import hashlib
import json
import os
import re
import time
from dataclasses import fields
from typing import Callable, Iterable, Union

from loguru import logger

from sweepai.config.server import (
    ANTHROPIC_API_KEY,
    AWS_ACCESS_KEY,
    AWS_SECRET_KEY,
    COHERE_API_KEY,
    GITHUB_APP_ID,
    GITHUB_APP_PEM,
    INSTALLATION_ID,
    LICENSE_KEY,
    OPENAI_API_KEY,
    VOYAGE_API_KEY,
)
from sweepai.o11y.event_logger import posthog

DEFAULT_BOT_SUFFIX = "\n\n*This is an automated message generated by [Sweep AI](https://sweep.dev).*"
BOT_SUFFIX = os.environ.get("BOT_SUFFIX", DEFAULT_BOT_SUFFIX).replace("\\n", "\n")

if BOT_SUFFIX != DEFAULT_BOT_SUFFIX:
    print(f"Using custom bot suffix: {BOT_SUFFIX}")

FASTER_MODEL_MESSAGE = """\
You ran out of the free tier GPT-4 tickets! Here are your options:
- You can get a free trial of Sweep Pro to get unlimited Sweep issues [here](https://buy.stripe.com/00g5npeT71H2gzCfZ8).
- You can run Sweep with your own Anthropic and OpenAI API keys [here](https://docs.sweep.dev/cli).
- You can book a chat with us set up Sweep Enterprise [here](https://calendly.com/d/2n5-3qf-9xy/user-interview).
"""

sep = "\n---\n"
bot_suffix_starring = ""
bot_suffix = f"""
{sep}
> [!TIP]
> To recreate the pull request, edit the issue title or description."""
discord_suffix = ""

stars_suffix = ""

collapsible_template = """
<details {opened}>
<summary>{summary}</summary>

{body}
</details>
"""

checkbox_template = "- [{check}] {filename}\n{instructions}\n"

num_of_snippets_to_query = 30
total_number_of_snippet_tokens = 15_000
num_full_files = 2


def ordinal(n):
    return str(n) + ("th" if 4 <= n <= 20 else {1: "st", 2: "nd", 3: "rd"}.get(n % 10, "th"))


def tail(logs: str, max_lines: int = 1000) -> str:
    logs_lines = logs.splitlines()
    if len(logs_lines) > max_lines:
        return "\n".join(logs_lines[-max_lines:])
    return logs


def truncate_text_width(text: str, width: int = 1000):
    lines = [line[:width] for line in text.splitlines()]
    return "\n".join(lines)


def create_collapsible(summary: str, body: str, opened: bool = False):
    return collapsible_template.format(summary=summary, body=body, opened="open" if opened else "")


def inline_code(text: str):
    return f"<code>{text}</code>" if text else ""


def code_block(text: str, xml: bool = False):
    if not text:
        return ""
    return f"```\n{text}\n```" if xml else f"<pre>\n{text}\n</pre>"


def blockquote(text: str):
    text = text.replace("\n•", "<br/>•")
    return f"<blockquote>\n{text}\n</blockquote>" if text else ""


def bold(text: str):
    return f"<b>{text}</b>" if text else ""


def markdown_list(items: list[str], bullet: str = "-"):
    bullet += " "
    return "\n".join([f"{bullet}{item}" for item in items])


def create_checkbox(title: str, body: str, checked: bool = False):
    return checkbox_template.format(check="X" if checked else " ", filename=title, instructions=body)


LANGUAGE_COMMENTS = {
    "py": "#",
    "js": "//",
    "ts": "//",
    "tsx": "//",
    "go": "//",
    "rb": "#",
    "java": "//",
    "c": "//",
    "cpp": "//",
    "cs": "//",
    "php": "//",
    "swift": "//",
    "kt": "//",
    "rs": "//",
    "pl": "#",
    "r": "#",
    "scala": "//",
    "lua": "--",
    "hs": "--",
    "sh": "#",
    "bash": "#",
}


def create_language_specific_comment(
    filename: str,
    comment_text: str,
    start_delimiter: str = "Pull Request Comment: ",
    end_delimiter: str = "",
):
    # Extract the file extension (or whole filename if no extension)
    ext = filename.split(".")[-1].lower()

    # Get the comment symbol for the extension
    comment_symbol = LANGUAGE_COMMENTS.get(ext, "#")  # Default to '#' if extension not found

    # Format the comment
    formatted_comment = f" {comment_symbol} {start_delimiter}{comment_text}{end_delimiter}"

    return formatted_comment


def format_comment_prompt_from_files(
    comment_path: str,
    comment_position: int,
    comment_diff_hunk: str,
    comment_body: str,
    comment_indicator: str = "▶",
):
    # delete the @@ -x1, y1 +x2, y2 @@ from the diff hunk by deleting until the second @@ in the code if it exists
    second_at_sign_index = comment_diff_hunk.find("@@ ", comment_diff_hunk.find("@@ ") + 1)
    if second_at_sign_index != -1:
        # this usually shouldn't delete any newlines
        comment_diff_hunk = comment_diff_hunk[second_at_sign_index + 3 :]
    diff_hunk_lines = comment_diff_hunk.split("\n")
    # ensure that we always add the comment somewhere
    for idx in range(len(diff_hunk_lines)):
        if idx == comment_position:
            diff_hunk_lines[idx] = (
                comment_indicator
                + " "
                + diff_hunk_lines[idx]
                + create_language_specific_comment(comment_path, comment_body.strip())
            )
            break
    else:
        diff_hunk_lines[-1] = (
            comment_indicator
            + " "
            + diff_hunk_lines[-1]
            + create_language_specific_comment(comment_path, comment_body.strip())
        )
    return ":\n" + "\n".join(diff_hunk_lines)


def strip_sweep(text: str):
    text = re.sub(r"^[Ss]weep\s?:", "", text).strip()
    text = re.sub(r"^@[Ss]weep(:?)", "", text).strip()
    return text


def clean_logs(logs: str):
    cleaned_logs = re.sub(r"\x1b\[.*?[@-~]", "", logs.replace("```", r"```"))
    cleaned_logs = re.sub("\n{2,}", "\n", cleaned_logs)
    cleaned_logs = re.sub("\r{2,}", "\n", cleaned_logs)
    cleaned_logs = cleaned_logs.strip("\n")
    cleaned_logs = cleaned_logs or "(nothing was outputted)"
    return cleaned_logs


def extract_lines(text: str, start: int, end: int):
    if not text:
        return ""
    lines = text.splitlines(keepends=True)
    return "".join(lines[max(0, start) : min(len(lines), end)])


def add_line_numbers(text: str, start: int = 0):
    if not text:
        return ""
    lines = text.splitlines(keepends=True)
    # Calculate the width needed for line numbers
    max_line_number = start + len(lines) - 1
    number_width = len(str(max_line_number))
    # Format each line with aligned numbers and separator
    numbered_lines = [f"{start + i:{number_width}} | {line}" for i, line in enumerate(lines)]
    return "".join(numbered_lines)


# start and end are inclusive
def remove_lines_from_text(text: str, start: int = 0, end: int = -1):
    if not text:
        return ""
    lines = text.splitlines(keepends=True)
    if end == -1:
        end = len(lines) - 1
    elif end >= len(lines):
        end = len(lines) - 1
    if start - 1 < 0:
        start = 1
    updated_lines = lines[: start - 1] + lines[end + 1 :]
    return "".join(updated_lines)


def to_branch_name(s, max_length=40):
    branch_name = s.strip().lower().replace(" ", "_")
    branch_name = re.sub(r"[^a-z0-9_]", "", branch_name)
    return branch_name[:max_length]


def get_hash():
    return hashlib.sha256(str(time.time()).encode()).hexdigest()[:10]


# used for getting all indices of a substring match
def get_all_indices_of_substring(content: str, substring: str):
    start = 0
    indices = []
    while True:
        index = content.find(substring, start)
        if index == -1:  # No more occurrences found
            break
        indices.append(index)
        start = index + 1  # Move past the last found occurrence
    return indices


def get_xml_parsing_regex(tag_name: str):
    return rf"<{tag_name}>(?P<{tag_name}>.*?)</{tag_name}>"


# converts a single arbitrary object to xml string format
def object_to_xml(object_: object, object_name: str, exclude_fields: list[str] = []):
    newline = "\n"
    if isinstance(object_, dict):
        object_fields = [
            f"<{field}>\n{object_[field].strip(newline)}\n</{field}>"
            for field in sorted(object_)
            if field not in exclude_fields
        ]
    elif isinstance(object_, str):
        object_fields = [f"<{object_name}>\n{object_.strip(newline)}\n</{object_name}>"]
    else:
        try:
            object_fields = [
                f"<{field.name}>\n{getattr(object_, field.name).strip(newline)}\n</{field.name}>"
                for field in fields(object_)
                if field.name not in exclude_fields
            ]
        except Exception as e:
            logger.error(f"Failed to convert object to xml: {e}")
            return ""
    fields_strings = "\n".join(object_fields)
    object_string = f"<{object_name}>\n{fields_strings}\n</{object_name}>"
    return object_string


def wrap_xml_tag(content: str, tag: str):
    stripped_content = content.strip("\n")
    return f"<{tag}>\n{stripped_content}\n</{tag}>"


# converts a list of objects to xml string format
def objects_to_xml(
    objects: list[object],
    object_name: str,
    outer_field_name: str = "",
    exclude_fields: list[str] = [],
):
    objects_string = ""
    for object in objects:
        objects_string += f"{object_to_xml(object, object_name, exclude_fields=exclude_fields)}\n"
    if outer_field_name:
        objects_string = f"<{outer_field_name}>\n{objects_string}</{outer_field_name}>"
    else:
        objects_string = f"<{object_name}s>\n{objects_string}</{object_name}s>"
    return objects_string


def extract_xml_tag(string: str, tag: str, include_closing_tag: bool = True):
    if not include_closing_tag:
        string += f"</{tag}>"
    pattern = rf"<{tag}>(.*?)</{tag}>"
    match_ = re.search(pattern, string, re.DOTALL)
    if match_ is None:
        return None
    return match_.group(1).strip("\n")

def extract_all_xml_tags(string: str, tag: str):
    pattern = rf"<{tag}>(.*?)</{tag}>"
    matches = re.findall(pattern, string, re.DOTALL)
    return [match.strip("\n") for match in matches]


def safe_extract_xml_tag(string: str, tag: str): # TODO: THIS IS A TYPO, supposed to be safe_extract_xml_tag
    """
    This one will never return None.
    """
    if f"<{tag}>" not in string:
        return string
    _before, content = string.split(f"<{tag}>", 1)
    if f"</{tag}>" not in content:
        return content.strip("\n")
    content, _after = content.split(f"</{tag}>", 1)
    return content.strip("\n")


def extract_object_fields_from_string(
    text: str, params: list[str]
) -> tuple[dict[str, str | dict[str, str]], bool, str]:
    """
    Extract a series of params from the text and return them as a dictionary.
    Example: extract_object_fields_from_string("<a>a</a><b>b</b>", ["a", "b"]) -> {a: a, b: b, raw_text: {a: "<a>a</a>", b: "<b>b</b>"}}
    """
    if not params:
        return {"content": text, "raw_text": {"content": text}}, False, ""
    object_args: dict[str, str | dict[str, str]] = {"raw_text": {}}
    failed = False
    failed_param = ""
    for param in params:
        regex = rf"<{param}>(?P<{param}>.*?)<\/{param}>"
        result = re.search(regex, text, re.DOTALL)
        try:
            object_args[param] = result.group(param).strip()  # type: ignore
            object_args["raw_text"][param] = result.group(0)  # type: ignore
        except AttributeError:
            failed = True
            failed_param = param
            break
    return object_args, failed, failed_param


def extract_objects_from_string(
    text: str,
    object_tag: str,
    object_params: list[str] = [],
    include_closing_tag: bool = True,
):
    """
    Extract xml objects from strings and returns them as a list of dictionaries.
    Example: extract_objects_from_string("<test><a>a</a><b>b</b></test>", "test", ["a","b"]) -> [{a: a, b: b, raw_text: {test: "<test><a>a</a><b>b</b></test>", a: "<a>a</a>", b: "<b>b</b>"}}]
    """
    extracted_objects = []
    object_regex = (
        rf"<{object_tag}>(?P<{object_tag}>.*?)<\/{object_tag}>"
        if include_closing_tag
        else rf"<{object_tag}>(?P<{object_tag}>.*?)(<\/{object_tag}>|\Z)"
    )
    object_matches = list(re.finditer(object_regex, text, re.DOTALL))
    failed_extraction = False
    for match in object_matches:
        full_raw_text = match.group(0)  # This includes the object_tag
        object_args, failed, failed_param = extract_object_fields_from_string(
            match.group(f"{object_tag}"), object_params
        )
        if not failed:
            object_args["raw_text"][object_tag] = full_raw_text  # type: ignore
            extracted_objects.append(object_args)
        else:
            failed_extraction = True
            posthog.capture(
                "extract_objects_from_string",
                "extract_objects_from_string failed",
                properties={
                    "failed_param": failed_param,
                    "text": text,
                    "object_tag": object_tag,
                },
            )
    return extracted_objects, failed_extraction


def rstrip_lines(text: str) -> str:
    """Claude likes to put trailing spaces at the end of lines. This function removes them."""
    return "\n".join([line.rstrip() for line in text.split("\n")])


def strip_triple_quotes(text: str) -> str:
    stripped_text = text.strip("\n").rstrip()
    if stripped_text.startswith("```"):
        stripped_text = stripped_text[stripped_text.find("\n") + 1 :]
    if stripped_text.endswith("```"):
        if "\n" in stripped_text:
            stripped_text = stripped_text[: stripped_text.rfind("\n")]
        else:
            stripped_text = stripped_text.rstrip("```")  # added this case to handle case where text = ```\n```, it previously returned only ``
    return stripped_text


def strip_triple_quotes_preserve_whitespace(text: str) -> str:
    """
    Strips triple quotes from the beginning and end of a string,
    preserving whitespace except when triple quotes are removed.
    """
    stripped_text = text

    # Check and remove leading triple quotes
    if stripped_text.lstrip().startswith("```"):
        stripped_text = stripped_text.lstrip()
        stripped_text = stripped_text[stripped_text.find("\n") + 1 :]

    # Check and remove trailing triple quotes
    if stripped_text.rstrip().endswith("```"):
        stripped_text = stripped_text.rstrip()
        stripped_text = stripped_text[: stripped_text.rfind("\n")]

    return stripped_text


def truncate_text_based_on_stop_sequence(text: str, stop_sequences: list[str]):
    for stop_sequence in stop_sequences:
        stop_index = text.find(stop_sequence)
        if stop_index != -1:
            text = text[: stop_index + len(stop_sequence)]
    return text


def extract_fenced_code_blocks(text: str) -> dict[str, dict[str, str]]:
    """
    Attempts to find all fenced code blocks in the given text and returns a dictionary of code blocks.\n
    Example:\n
    Input:\n
    ```python
    print("test")
    ```
    Output:
    {
        'print("test")': {
            'full_text': '```python\nprint("test")\n```',
            'language': 'python'
        }
    }
    """
    pattern = re.compile(r"```(?P<language>\w+)\n(?P<code>[\s\S]*?)```")
    code_blocks = {}

    # Find all matches
    matches = pattern.finditer(text)

    # Iterate over the matches and populate the code_blocks dictionary
    for match in matches:
        language = match.group("language")
        code = match.group("code")
        full_text = match.group(0)  # This captures the entire matched text including backticks and language

        # Use the actual code as the key, and the value is a dict containing full text and language
        code_blocks[code] = {"full_text": full_text, "language": language}

    return code_blocks


def find_filenames(text: str):
    pattern = r"(?:^|\b)((?:\.{1,2}/)?(?:[\w-]+/)*[\w-]+\.[\w-]+)(?=$|\b)"
    return list(dict.fromkeys(re.findall(pattern, text)))


def ordered_dedup(lst: list[str], key: Callable = lambda x: x):
    result = []
    seen = set()
    for item in lst:
        if key(item) not in seen:
            seen.add(key(item))
            result.append(item)
    return result


def pack_items_for_prompt(
    iterable: Iterable,
    string_function: Union[callable, None],
    token_limit: int,
    char_token_ratio: int = 3.5,
) -> list:
    """
    Packs items from an iterable into a list of strings, using a string function to convert each item to a string.
    The total number of tokens in the packed items will not exceed the token limit.
    """
    char_limit = token_limit * char_token_ratio
    packed_items = []
    current_str = ""
    for item in iterable:
        item_str = string_function(item) if string_function else str(item)
        if len(current_str) + len(item_str) <= char_limit:
            packed_items.append(item)
            current_str += item_str
        else:
            break
    logger.info(
        f"Removed {len(iterable) - len(packed_items)} items to fit within the token limit. Final token estimate: {int(len(current_str) // char_token_ratio)}"
    )
    return packed_items


def format_json_string(jsonifiable: str) -> str:
    """
    format a json string input
    """
    try:
        return json.dumps(json.loads(jsonifiable), indent=4)
    except Exception as e:
        logger.error(f"Failed to pretty print JSONifiable object: {e}")
        return str(jsonifiable)


def get_indentation(text: str) -> str:
    """
    Get the indentation depth of a string.
    """
    indent = text.removesuffix(text.lstrip())
    return " " if "\t" not in indent else "\t", len(indent)


def dedent_text(text: str, indent: str) -> str:
    """
    Dedent the text by the amount of leading spaces.
    """
    return text.removesuffix(indent)


def find_match_line_index(text: str, pattern: str):
    """
    Find the first occurrence of a multiline pattern in a multiline text. Returns the line index where the pattern starts. (this is claude generated)
    """
    text_lines = text.splitlines()
    pattern_lines = pattern.splitlines()

    if not pattern_lines or len(pattern_lines) > len(text_lines):
        return -1

    # special case for single line patterns
    if len(pattern_lines) == 1:
        for i, line in enumerate(text_lines):
            if pattern in line:
                return i
        return -1

    # added indentation support
    indentation_type, indentation_depth = get_indentation(pattern_lines[0])
    dedented_pattern_lines = pattern_lines

    # This is inefficient but it works for now

    # loop over all dedentations of the pattern
    for iter in range(indentation_depth + 1):
        if iter > 0:
            dedented_pattern_lines = [line.removeprefix(indentation_type) for line in dedented_pattern_lines]
        dedented_first_pattern_line = dedented_pattern_lines[0]
        rstripped_first_pattern_line = dedented_first_pattern_line.rstrip()
        last_pattern_line = dedented_pattern_lines[-1]
        rstripped_last_pattern_line = last_pattern_line.rstrip()

        for i, line in enumerate(text_lines):
            rstripped_line = line.rstrip()
            if rstripped_line.endswith(rstripped_first_pattern_line):
                indent = rstripped_line.removesuffix(rstripped_first_pattern_line)

                if len(dedented_pattern_lines) == 1:
                    return i

                # Check if there are enough remaining lines
                if i + len(dedented_pattern_lines) > len(text_lines):
                    continue

                match = True
                for j in range(1, len(dedented_pattern_lines) - 1):
                    if text_lines[i + j].rstrip() != indent + dedented_pattern_lines[j].rstrip():
                        match = False
                        break

                # TODO: ensure proper indentation as well
                if match and text_lines[i + len(dedented_pattern_lines) - 1].startswith(
                    indent + rstripped_last_pattern_line
                ):
                    return i

    return -1  # Pattern not found


def find_match_with_leading_spaces(text: str, pattern: str) -> tuple[int, int]:
    """
    Find the first occurrence of a multiline pattern in a multiline text,
    allowing for any number of leading spaces in each line of the pattern.
    Returns the line index where the pattern starts, or -1 if not found.
    Also returns if the number of leading spaces required to match.
    """
    pattern_lines = pattern.splitlines()

    if not pattern_lines:
        return -1, 0

    # check for up to 20 leading tabs assuming a tab is 2 spaces
    max_spaces = 20
    space_multiplier = 2

    for spaces in range(max_spaces + 1):
        spaced_pattern = "\n".join(" " * spaces * space_multiplier + line for line in pattern_lines)
        result = find_match_line_index(text, spaced_pattern)
        if result != -1:
            return result, spaces * space_multiplier

    return -1, 0


def reset_rightmost_string(text: str, old: str, new: str):
    """
    Resets the right most instance of a substring in a string.
    """
    return text.rsplit(old, 1)[0] + new


def should_truncate(file_contents: str, partial_contents: str) -> bool:
    """
    Return a bool for whether partial_contents is not in file_contents,
    in which case it should be truncated.
    """
    if not partial_contents:
        return False
    return partial_contents not in file_contents


def truncate(file_contents, partial_contents):
    """
    Find largest prefix of partial_contents that is a substring of file_contents.
    Ignores indentation errors.
    """

    lines = file_contents.splitlines()
    partial_contents_lines = partial_contents.splitlines()

    for i in range(len(partial_contents_lines), 0, -1):
        partial_contents_prefix = "\n".join(partial_contents_lines[:i])
        matching_line_index, _ = find_match_with_leading_spaces(file_contents, partial_contents_prefix)
        if matching_line_index != -1:
            return "\n".join(lines[matching_line_index:matching_line_index + i])
    return ""


def wildcard_match(username: str, wildcard_list: list[str]) -> bool:
    if username in wildcard_list:
        return True
    if len(wildcard_list) == 1 and wildcard_list[0] == "*":
        return True
    return False


def truncate_commit_message(original_message: str, length: int = 64) -> str:
    safe_length_message = (
        original_message
        if len(original_message) < length
        else original_message[: original_message.rfind(" ", 0, length)]
    )
    if len(safe_length_message) >= length:
        safe_length_message = safe_length_message[:length]
    safe_length_message = safe_length_message.rstrip().rstrip(",")
    assert len(safe_length_message) <= length, "Commit message too long"
    return safe_length_message


def sanitize_string_for_github(message: str):
    secrets = [
        GITHUB_APP_PEM,
        GITHUB_APP_ID,
        ANTHROPIC_API_KEY,
        OPENAI_API_KEY,
        INSTALLATION_ID,
        os.environ.get("GITHUB_PAT", ""),
        COHERE_API_KEY,
        LICENSE_KEY,
        VOYAGE_API_KEY,
        AWS_ACCESS_KEY,
        AWS_SECRET_KEY,
    ]
    secrets = [secret for secret in secrets if secret]
    for secret in secrets:
        if secret and secret in message:
            message = message.replace(secret, "*" * len(secret))
    return message
