import sys
import os
import ast
import textwrap
import re
# Add the pipeline/ directory to sys.path
script_dir = os.path.dirname(os.path.abspath(__file__))
pipeline_root = os.path.abspath(os.path.join(script_dir, ".."))
sys.path.insert(0, pipeline_root)

# from app.search.search_utils import *

def remove_line_numbers(function_code: str) -> str:
    """
    Improved version of line number removal and indentation normalization.
    """
    lines = function_code.splitlines()
    line_number_pattern = re.compile(r'^\s*\d+\s')

    stripped_lines = []
    for line in lines:
        if line_number_pattern.match(line):
            line = line_number_pattern.sub('', line)
        stripped_lines.append(line)

    non_empty = [l for l in stripped_lines if l.strip()]
    if not non_empty:
        return ""

    min_indent = min(len(l) - len(l.lstrip()) for l in non_empty)
    normalized_lines = [(l[min_indent:] if l.strip() else "") for l in stripped_lines]
    return "\n".join(normalized_lines)

def make_ast_parsable(function_code: str) -> ast.AST | None:
    """
    Try to ensure the given function code is parsable by the Python AST module.
    Returns the AST tree if successful, otherwise returns None.
    """
    code = remove_line_numbers(function_code)
    code = textwrap.dedent(code)

    try:
        return ast.parse(code)
    except SyntaxError:
        # Try wrapping in a dummy class or module
        try:
            wrapped = f"def _dummy_wrapper():\n{textwrap.indent(code, '    ')}"
            return ast.parse(wrapped)
        except SyntaxError as e:
            print(f"[AST Error] Still not parsable: {e}")
            return None

exp_func = '''    def test_rtol_diff(self, capsys):
        a = np.arange(100, dtype=float).reshape(10, 10)
        hdu_a = PrimaryHDU(data=a)
        b = a.copy()
        b[1, 0] = 11
        hdu_b = PrimaryHDU(data=b)
        tmp_a = self.temp("testa.fits")
        tmp_b = self.temp("testb.fits")
        hdu_a.writeto(tmp_a)
        hdu_b.writeto(tmp_b)
        numdiff = fitsdiff.main(["-r", "1e-2", tmp_a, tmp_b])
        assert numdiff == 1
        out, err = capsys.readouterr()
        assert (
            out
            == f"""
 fitsdiff: {version}
 a: {tmp_a}
 b: {tmp_b}
 Maximum number of different data values to be reported: 10
 Relative tolerance: 0.01, Absolute tolerance: 0.0

Primary HDU:

   Data contains differences:
     Data differs at [1, 2]:
        a> 10.0
         ?  ^
        b> 11.0
         ?  ^
different pixels found (1.00% different).
"""
        )
        assert err == ""'''
        
        
normalized_code = remove_line_numbers(exp_func)
print(normalized_code)
tree = make_ast_parsable(normalized_code)
print("Parsable!" if tree else "Parsing failed.")
