import re
import ast
import os

def file_read(filepath):
    """
    Read a Python function from a file as text.
    
    Args:
        filepath: Path to the Python file containing the function
        
    Returns:
        str: The contents of the file as a string
    """
    with open(filepath, 'r', encoding='utf-8') as f:
        return f.read()
    
def read_lines(filepath):
    with open(filepath, 'r') as file:
        lines = file.readlines()
        
    return lines

def write_lines(filepath, lines):
    with open(filepath, 'w') as file:
        file.writelines(lines)

def write_lines2(filepath, lines):
    with open(filepath, 'w') as file:
        file.writelines(line if line.endswith('\n') else line + '\n' for line in lines)
    
def file_write(content, filepath):
    """
    Write a string to a file.
    
    Args:
        filepath: Path to the file to write to
        content: String content to write
    """
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(content)
        

def find_import_idx(lines):
    """
    Find indices of lines that are import statements and come before any function definitions.
    
    Args:
        lines: List of strings to search
        
    Returns:
        List of indices of import statements before any def/async def
    """
    import_indices = []
    
    for i, line in enumerate(lines):
        stripped = line.lstrip()
        
        # Check if we've hit a function definition
        if stripped.startswith('def ') or stripped.startswith('async def '):
            break
            
        # Check if it's an import statement
        if stripped.startswith('import ') or (stripped.startswith('from ') and 'import' in stripped):
            import_indices.append(i)
    
    max_idx = max(import_indices)
    
    return max_idx


def split_by_newlines(s):
    result = []
    current = []
    in_single = False
    in_double = False
    in_triple_single = False
    in_triple_double = False
    i = 0
    
    while i < len(s):
        # Check for triple quotes first
        if i + 2 < len(s):
            if s[i:i+3] == '"""':
                if in_triple_double:
                    in_triple_double = False
                elif not in_single and not in_double and not in_triple_single:
                    in_triple_double = True
                current.append(s[i:i+3])
                i += 3
                continue
            elif s[i:i+3] == "'''":
                if in_triple_single:
                    in_triple_single = False
                elif not in_single and not in_double and not in_triple_double:
                    in_triple_single = True
                current.append(s[i:i+3])
                i += 3
                continue
        
        # Check for single quotes
        if s[i] == "'" and not in_double and not in_triple_double and not in_triple_single:
            in_single = not in_single
            current.append(s[i])
        # Check for double quotes
        elif s[i] == '"' and not in_single and not in_triple_single and not in_triple_double:
            in_double = not in_double
            current.append(s[i])
        # Check for newline
        elif s[i] == '\n':
            if in_single or in_double or in_triple_single or in_triple_double:
                current.append(s[i])
            else:
                result.append(''.join(current))
                current = []
        else:
            current.append(s[i])
        
        i += 1
    
    # Add remaining content
    if current:
        result.append(''.join(current))
    
    return result

def add_code_me(string, filepath):
    lines = read_lines(filepath)
    import_idx = find_import_idx(lines)
    string_lines = split_by_newlines(string)
    string_import_idx = find_import_idx(string_lines)
    
    first_file_lines = lines[:import_idx]
    other_file_lines = lines[(import_idx+1):]
    
    first_string_lines = string_lines[:string_import_idx]
    other_string_lines = string_lines[(string_import_idx+1):]
    
    first_lines = first_file_lines + first_string_lines
    other_lines = other_file_lines + other_string_lines
    
    print(first_lines)
    print(other_lines)
    
    write_lines(filepath, lines)
    
def add_code(code: str, filename: str) -> None:
    # Parse the code string to extract lines while respecting quoted strings
    lines = []
    current_line = []
    in_single_quote = False
    in_double_quote = False
    in_triple_single = False
    in_triple_double = False
    i = 0

    while i < len(code):
        char = code[i]

        # Check for triple quotes first
        if i + 2 < len(code):
            three_chars = code[i:i+3]
            if three_chars == '"""' and not in_single_quote and not in_triple_single:
                current_line.append(three_chars)
                in_triple_double = not in_triple_double
                i += 3
                continue
            elif three_chars == "'''" and not in_double_quote and not in_triple_double:
                current_line.append(three_chars)
                in_triple_single = not in_triple_single
                i += 3
                continue

        # Check for escape sequences
        if char == '\\' and i + 1 < len(code):
            current_line.append(char)
            current_line.append(code[i + 1])
            i += 2
            continue

        # Check for single/double quotes (only if not in triple quotes)
        if not in_triple_single and not in_triple_double:
            if char == '"' and not in_single_quote:
                in_double_quote = not in_double_quote
            elif char == "'" and not in_double_quote:
                in_single_quote = not in_single_quote

        # Check for newline
        if char == '\n' and not in_single_quote and not in_double_quote and not in_triple_single and not in_triple_double:
            lines.append(''.join(current_line))
            current_line = []
        else:
            current_line.append(char)

        i += 1

    # Add any remaining content as the last line
    if current_line:
        lines.append(''.join(current_line))

    # Separate import lines from other lines
    import_lines = []
    other_lines = []

    for line in lines:
        stripped = line.strip()
        if stripped.startswith('import ') or (stripped.startswith('from ') and ' import ' in stripped):
            import_lines.append(line)
        else:
            other_lines.append(line)

    # Read existing file content
    try:
        with open(filename, 'r') as f:
            existing_content = f.read()
    except FileNotFoundError:
        existing_content = ""

    # Parse existing file to find where imports end
    existing_lines = existing_content.split('\n') if existing_content else []

    # Find existing imports and their positions
    existing_imports = []
    last_import_index = -1

    in_single_quote = False
    in_double_quote = False
    in_triple_single = False
    in_triple_double = False

    for idx, line in enumerate(existing_lines):
        # Track quote state across lines
        i = 0
        while i < len(line):
            char = line[i]

            if i + 2 < len(line):
                three_chars = line[i:i+3]
                if three_chars == '"""' and not in_single_quote and not in_triple_single:
                    in_triple_double = not in_triple_double
                    i += 3
                    continue
                elif three_chars == "'''" and not in_double_quote and not in_triple_double:
                    in_triple_single = not in_triple_single
                    i += 3
                    continue

            if char == '\\' and i + 1 < len(line):
                i += 2
                continue

            if not in_triple_single and not in_triple_double:
                if char == '"' and not in_single_quote:
                    in_double_quote = not in_double_quote
                elif char == "'" and not in_double_quote:
                    in_single_quote = not in_single_quote

            i += 1

        # Only check for imports if we're not inside a string
        if not in_single_quote and not in_double_quote and not in_triple_single and not in_triple_double:
            stripped = line.strip()
            if stripped.startswith('import ') or (stripped.startswith('from ') and ' import ' in stripped):
                existing_imports.append(line)
                last_import_index = idx

    # Filter out duplicate imports
    new_imports = []
    for imp in import_lines:
        if imp not in existing_imports and imp not in new_imports:
            new_imports.append(imp)

    # Build the new file content
    if last_import_index >= 0:
        # Insert new imports after the last existing import
        new_content_lines = existing_lines[:last_import_index + 1] + new_imports + existing_lines[last_import_index + 1:]
    elif new_imports:
        # No existing imports, add new imports at the top
        new_content_lines = new_imports + existing_lines
    else:
        new_content_lines = existing_lines

    # Add the non-import code at the bottom
    if other_lines:
        # Join with actual newlines for the new code
        new_code_str = '\n'.join(other_lines)
        if new_content_lines and new_content_lines[-1] != '':
            new_content_lines.append('')
        new_content_lines.append(new_code_str)

    # Write back to file
    with open(filename, 'w') as f:
        f.write('\n'.join(new_content_lines))

def split2functions(filename):
    with open(filename, 'r') as f:
        content = f.read()
    
    tree = ast.parse(content)
    lines = content.splitlines(keepends=True)
    
    result = []
    
    # Find all top-level function definitions
    functions = [node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)]
    # Filter to only top-level functions
    top_level_functions = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
    
    if not top_level_functions:
        # No functions, return everything as preamble
        return [content]
    
    # Get preamble (everything before first function)
    first_func = top_level_functions[0]
    first_func_line = first_func.lineno
    preamble = ''.join(lines[:first_func_line - 1])
    result.append(preamble)
    
    # Get each function
    for i, func in enumerate(top_level_functions):
        start_line = func.lineno - 1
        if i + 1 < len(top_level_functions):
            end_line = top_level_functions[i + 1].lineno - 1
        else:
            end_line = len(lines)
        
        func_text = ''.join(lines[start_line:end_line])
        result.append(func_text)
    
    return result

def edit(filepath, line=None):
    """
    Open a Python file in Spyder's editor, optionally at a specific line.

    Parameters
    ----------
    filepath : str
        Path to the file to open in Spyder's editor.
    line : int, optional
        Line number to jump to when opening the file.

    Examples
    --------
    >>> spyder_edit('myfile.py')
    >>> spyder_edit('myfile.py', line=42)
    """
    if line is not None:
        get_ipython().run_line_magic('edit', f'{filepath}:{line}')
    else:
        get_ipython().run_line_magic('edit', filepath)