import sys, traceback, inspect, linecache
from fansi_text import Yellow, Bold, Blue, Underline, Green, BrightMagenta, BrightYellow

def _get_context(filename, line_num):
    # Get all lines from file
    all_lines = linecache.getlines(filename)
    target_index = line_num - 1  # zero-based index

    # Collect context: 2 non-empty before and after
    context = []
    # Backward
    before = []
    i = target_index - 1
    while i >= 0 and len(before) < 2:
        if all_lines[i].strip():
            before.insert(0, (i + 1, all_lines[i].rstrip()))
        i -= 1
    # Target line
    context.append((line_num, all_lines[target_index].rstrip()))
    # Forward
    after = []
    i = target_index + 1
    while i < len(all_lines) and len(after) < 2:
        if all_lines[i].strip():
            after.append((i + 1, all_lines[i].rstrip()))
        i += 1

    # Combine
    return before + context + after

def warn(message: str, warn_type: str = "Warning") -> None:
    """Print a formatted warning with traceback and yellow text.

    Args:
        message (str): The main warning text to display.
        warn_type (str, optional): An optional custom name for the warning; displayed in bold. Defaults to "Warning".
    """
    exc_type, exc_value, exc_traceback = sys.exc_info()
    
    if exc_traceback:
        frame = exc_traceback.tb_frame
        line_num = traceback.extract_tb(exc_traceback)[0].lineno
    else:
        frame = inspect.currentframe().f_back
        line_num = frame.f_lineno
        
    filename = frame.f_code.co_filename
    
    # Get 2 lines before and after
    context_lines = _get_context(filename, line_num)

    print(f'  File ' + BrightMagenta(filename) + ', line ' + BrightMagenta(frame.f_lineno))
    
    spacer = "  "
    for i, line in context_lines:
        if i in [10, 100, 1000, 10000]:
            spacer = spacer[:-1]
            
        # print(f"Spacer len = {len(spacer)}")
        if i == line_num:
            print(spacer + Yellow(Bold("---> ")) + BrightYellow(Underline(f"{i}")) + f"| {line.rstrip()}")
        else:
            print(spacer + "     " + Green(f"{i}") + f"| {line}")

    print(Yellow(Bold(warn_type)) + ': ' + Yellow(message))
