import re

from pr.config import LANGUAGE_KEYWORDS
from pr.ui.colors import Colors


def highlight_code(code, language=None, syntax_highlighting=True):
    if not syntax_highlighting:
        return code

    if not language:
        if "def " in code or "import " in code:
            language = "python"
        elif "function " in code or "const " in code:
            language = "javascript"
        elif "public " in code or "class " in code:
            language = "java"

    if language and language in LANGUAGE_KEYWORDS:
        keywords = LANGUAGE_KEYWORDS[language]
        for keyword in keywords:
            pattern = r"\b" + re.escape(keyword) + r"\b"
            code = re.sub(pattern, f"{Colors.BLUE}{keyword}{Colors.RESET}", code)

    code = re.sub(r'"([^"]*)"', f'{Colors.GREEN}"\\1"{Colors.RESET}', code)
    code = re.sub(r"'([^']*)'", f"{Colors.GREEN}'\\1'{Colors.RESET}", code)

    code = re.sub(r"#(.*)$", f"{Colors.GRAY}#\\1{Colors.RESET}", code, flags=re.MULTILINE)
    code = re.sub(r"//(.*)$", f"{Colors.GRAY}//\\1{Colors.RESET}", code, flags=re.MULTILINE)

    return code


def render_markdown(text, syntax_highlighting=True):
    if not syntax_highlighting:
        return text

    code_blocks = []

    def extract_code_block(match):
        lang = match.group(1) or ""
        code = match.group(2)
        highlighted_code = highlight_code(code.strip("\n"), lang, syntax_highlighting)
        placeholder = f"%%CODEBLOCK{len(code_blocks)}%%"
        full_block = f"{Colors.GRAY}```{lang}{Colors.RESET}\n{highlighted_code}\n{Colors.GRAY}```{Colors.RESET}"
        code_blocks.append(full_block)
        return placeholder

    text = re.sub(r"```(\w*)\n(.*?)\n?```", extract_code_block, text, flags=re.DOTALL)

    inline_codes = []

    def extract_inline_code(match):
        code = match.group(1)
        placeholder = f"%%INLINECODE{len(inline_codes)}%%"
        inline_codes.append(f"{Colors.YELLOW}{code}{Colors.RESET}")
        return placeholder

    text = re.sub(r"`([^`]+)`", extract_inline_code, text)

    lines = text.split("\n")
    processed_lines = []
    for line in lines:
        if line.startswith("### "):
            line = f"{Colors.BOLD}{Colors.GREEN}{line[4:]}{Colors.RESET}"
        elif line.startswith("## "):
            line = f"{Colors.BOLD}{Colors.BLUE}{line[3:]}{Colors.RESET}"
        elif line.startswith("# "):
            line = f"{Colors.BOLD}{Colors.MAGENTA}{line[2:]}{Colors.RESET}"
        elif line.startswith("> "):
            line = f"{Colors.CYAN}> {line[2:]}{Colors.RESET}"
        elif re.match(r"^\s*[\*\-\+]\s", line):
            match = re.match(r"^(\s*)([\*\-\+])(\s+.*)", line)
            if match:
                line = (
                    f"{match.group(1)}{Colors.YELLOW}{match.group(2)}{Colors.RESET}{match.group(3)}"
                )
        elif re.match(r"^\s*\d+\.\s", line):
            match = re.match(r"^(\s*)(\d+\.)(\s+.*)", line)
            if match:
                line = (
                    f"{match.group(1)}{Colors.YELLOW}{match.group(2)}{Colors.RESET}{match.group(3)}"
                )
        processed_lines.append(line)
    text = "\n".join(processed_lines)

    text = re.sub(
        r"\[(.*?)\]\((.*?)\)",
        f"{Colors.BLUE}\\1{Colors.RESET}{Colors.GRAY}(\\2){Colors.RESET}",
        text,
    )
    text = re.sub(r"~~(.*?)~~", f"{Colors.GRAY}\\1{Colors.RESET}", text)
    text = re.sub(r"\*\*(.*?)\*\*", f"{Colors.BOLD}\\1{Colors.RESET}", text)
    text = re.sub(r"__(.*?)__", f"{Colors.BOLD}\\1{Colors.RESET}", text)
    text = re.sub(r"\*(.*?)\*", f"{Colors.CYAN}\\1{Colors.RESET}", text)
    text = re.sub(r"_(.*?)_", f"{Colors.CYAN}\\1{Colors.RESET}", text)

    for i, code in enumerate(inline_codes):
        text = text.replace(f"%%INLINECODE{i}%%", code)
    for i, block in enumerate(code_blocks):
        text = text.replace(f"%%CODEBLOCK{i}%%", block)

    return text