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

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