#!/usr/bin/env python3
"""Validate Python bindings against C test files.

Parses C test files to extract inputs and expected outputs, then
calls Python bindings and compares. For test files that can't be
parsed, falls back to calling C via subprocess.
"""

import ctypes
import importlib.util
import inspect
import math
import os
import re
import sys

import numpy as np

_ts_spec = importlib.util.spec_from_file_location(
    "ts",
    os.path.join(os.path.dirname(__file__), "..", "python", "ts", "__init__.py"),
)
if _ts_spec and _ts_spec.loader:
    ts = importlib.util.module_from_spec(_ts_spec)
    _ts_spec.loader.exec_module(ts)
else:
    print("Error: could not load ts module", file=sys.stderr)
    sys.exit(1)

_parser_spec = importlib.util.spec_from_file_location(
    "ts_parser",
    os.path.join(os.path.dirname(__file__), "..", "scripts", "ts_parser.py"),
)
if _parser_spec and _parser_spec.loader:
    _ts_parser_mod = importlib.util.module_from_spec(_parser_spec)
    _parser_spec.loader.exec_module(_ts_parser_mod)
    parse_names_only = _ts_parser_mod.parse_names_only
else:
    print("Error: could not load ts_parser", file=sys.stderr)
    sys.exit(1)


def get_default_params(fname, x_n):
    """Extract default params from a Python binding's signature.

    Returns a dict of param name -> value for non-array parameters,
    or None if the function has no extra parameters.
    """
    func = getattr(ts, fname, None)
    if func is None:
        return None
    sig = inspect.signature(func)
    params = {}
    for pname, param in sig.parameters.items():
        if param.annotation is np.ndarray:
            continue
        if param.default is not inspect.Parameter.empty:
            params[pname] = param.default
    return params if params else None


LIB_PATH = os.path.join(os.path.dirname(__file__), "..", "build", "libts.so")
N = 100
SEED = 42


def load_func_meta():
    return parse_names_only()


def _safe_float(text):
    """Safely convert a C expression to float without eval()."""
    text = text.strip()
    if not text:
        return None
    text = re.sub(r"\bM_PI\b", "3.14159265358979323846", text)
    text = re.sub(r"\bM_E\b", "2.71828182845904523536", text)
    text = re.sub(r"\bPI\b", "3.14159265358979323846", text)
    # Only allow simple numeric expressions: numbers, operators, parens, abs/min/max
    if not re.match(r"^[\d\s\+\-\*\/\(\)\.,eE]+$", text):
        return None
    try:
        import ast

        tree = ast.parse(text, mode="eval")
        # Only allow safe node types
        for node in ast.walk(tree):
            if not isinstance(
                node,
                (
                    ast.Expression,
                    ast.Constant,
                    ast.BinOp,
                    ast.UnaryOp,
                    ast.USub,
                    ast.UAdd,
                    ast.Add,
                    ast.Sub,
                    ast.Mult,
                    ast.Div,
                ),
            ):
                return None
        return float(
            eval(text, {"__builtins__": {}, "abs": abs, "min": min, "max": max})
        )
    except Exception:
        return None


def parse_array_initializer(text):
    """Parse a C array initializer like { 1, 2, 3.5, -4.0 } into a list of floats."""
    m = re.search(r"\{\s*(.*?)\}\s*;", text, re.DOTALL)
    if not m:
        return None
    values = []
    for v in m.group(1).split(","):
        v = v.strip().split("//")[0].strip()
        if not v:
            continue
        result = _safe_float(v)
        if result is None:
            return None
        values.append(result)
    return values if values else None


def parse_malloc_loops(content):
    """Parse malloc + loop patterns to compute inputs.

    Handles patterns like:
        uint_t n = 14;
        fp_t *x[4];
        x[0] = malloc( n * sizeof( fp_t ) );
        ...
        for ( uint_t i = 0; i < n; ++i )
        {
            x[0][i] = 10 + i;
            x[1][i] = 12 + i;
        }

    Returns (inputs, n) or (None, None).
    """
    # Find n
    n_match = re.search(r"uint_t\s+n\s*=\s*(\d+)", content)
    if not n_match:
        return None, None
    n = int(n_match.group(1))

    # Find malloc declarations
    malloc_match = re.search(r"fp_t\s+\*\s*x\[(\d+)\]", content)
    if not malloc_match:
        return None, None
    x_n = int(malloc_match.group(1))

    # Find loop body: x[i][j] = expr;
    loop_match = re.search(
        r"for\s*\([^)]*i\s*=\s*0[^)]*i\s*<\s*n[^)]*\+\+i\s*\)\s*\{([^}]+)\}",
        content,
        re.DOTALL,
    )
    if not loop_match:
        return None, None

    loop_body = loop_match.group(1)
    formulas = {}
    for fm in re.finditer(r"x\[(\d+)\]\[i\]\s*=\s*([^;]+);", loop_body):
        idx = int(fm.group(1))
        expr = fm.group(2).strip()
        formulas[idx] = expr

    if not formulas:
        return None, None

    # Compute inputs
    inputs = {}
    for idx, expr in formulas.items():
        arr = []
        for i in range(n):
            try:
                safe_expr = expr.replace("i", str(i)).replace("n", str(n))
                if not re.match(r"^[\d\s\+\-\*\/\(\)\.,]+$", safe_expr.strip()):
                    return None, None
                val = float(
                    eval(
                        safe_expr,
                        {"__builtins__": {}, "abs": abs, "min": min, "max": max},
                    )
                )
                arr.append(val)
            except Exception:
                return None, None
        inputs[idx] = arr

    # Build ordered list
    result = []
    for i in range(x_n):
        if i in inputs:
            result.append(inputs[i])
        else:
            return None, None

    return result, n


def parse_test_file(test_path):
    """Parse a C test file to extract inputs, expected outputs, and constants.

    Returns dict with keys: inputs (list of lists), expected (list of lists),
    constants (list of floats), n (int), or None if parsing fails.
    """
    with open(test_path, encoding="utf-8") as f:
        content = f.read()

    # Try malloc + loop pattern first
    malloc_inputs, malloc_n = parse_malloc_loops(content)
    if malloc_inputs:
        # Find expected arrays
        all_arrays = {}
        for m in re.finditer(
            r"fp_t\s+(?:const\s+)?(\w+)\[\s*\]\s*=\s*\{([^}]*)\}\s*;",
            content,
            re.DOTALL,
        ):
            arr_name = m.group(1)
            values = parse_array_initializer(m.group(0))
            if values is not None:
                all_arrays[arr_name] = values

        # Find expected pointer array
        exp_match = re.search(
            r"fp_t\s+(?:\*\s*)?const\s*\*\s*const\s+e\w*\[\s*\d*\s*\]\s*=\s*\{([^}]*)\}",
            content,
        )
        expected = []
        if exp_match:
            exp_names = [n.strip() for n in exp_match.group(1).split(",") if n.strip()]
            for name in exp_names:
                if name in all_arrays:
                    expected.append(all_arrays[name])

        # Find params
        params = None
        param_match = re.search(r"(?:c_t|params_t)\s+c\s*=\s*\{([^}]*)\}", content)
        if param_match:
            params = {}
            for fm in re.finditer(r"\.(\w+)\s*=\s*([^,}]+)", param_match.group(1)):
                fname = fm.group(1)
                fval_text = fm.group(2).strip()
                result = _safe_float(fval_text)
                if result is not None:
                    params[fname] = result

        # Also look for scalar parameter: uint_t k = 3; fp_t c = 2; fp_t threshold = 0.5;
        if params is None:
            # Try to extract params from ts_test call's 6th argument
            ts_test_match = re.search(
                r"ts_test\s*\([^,]+,\s*[^,]+,\s*[^,]+,\s*[^,]+,\s*[^,]+,\s*(&?(\w+)|NULL)",
                content,
            )
            if ts_test_match:
                param_arg = ts_test_match.group(2)
                if param_arg and param_arg != "NULL":
                    # Check if it's an array like "constants[N] = { a, b }"
                    arr_match = re.search(
                        rf"fp_t\s+{param_arg}\s*\[\s*\d+\s*\]\s*=\s*\{{([^}}]*)\}}",
                        content,
                    )
                    if arr_match:
                        # Resolve variable names to values
                        var_names = [v.strip() for v in arr_match.group(1).split(",")]
                        resolved = {}
                        for vn in var_names:
                            vn = vn.strip()
                            if not vn:
                                continue
                            # Try direct numeric value
                            result = _safe_float(vn)
                            if result is not None:
                                resolved[vn] = result
                                continue
                            # Try to find variable declaration
                            var_match = re.search(
                                rf"(?:uint_t|fp_t|int)\s+(?:const\s+)?{vn}\s*=\s*([^;]+);",
                                content,
                            )
                            if var_match:
                                result = _safe_float(var_match.group(1).strip())
                                if result is not None:
                                    resolved[vn] = result
                        if resolved:
                            params = resolved

            if params is None:
                scalar_match = re.search(
                    r"(?:uint_t|fp_t|int)\s+(?:const\s+)?(\w+)\s*=\s*([^;]+);", content
                )
                if scalar_match:
                    var_name = scalar_match.group(1)
                    # Skip common non-param variables
                    if var_name not in ("n", "offset", "y_n", "i", "j"):
                        val_text = scalar_match.group(2).strip()
                        result = _safe_float(val_text)
                        if result is not None:
                            params = {var_name: result}

        return {
            "inputs": malloc_inputs,
            "expected": expected,
            "constants": [],
            "params": params,
            "n": malloc_n,
        }

    # Find all array declarations: fp_t const name[] = { ... };
    all_arrays = {}
    for m in re.finditer(
        r"fp_t\s+(?:const\s+)?(\w+)\[\s*\]\s*=\s*\{([^}]*)\}\s*;", content, re.DOTALL
    ):
        arr_name = m.group(1)
        values = parse_array_initializer(m.group(0))
        if values is not None:
            all_arrays[arr_name] = values

    if not all_arrays:
        return None

    # Find input pointer array: fp_t const * const inp[N] = { x0, x1, ... };
    inp_match = re.search(
        r"fp_t\s+const\s*\*\s*const\s+(?:inp\w*|x\w*)\[\s*\d*\]\s*=\s*\{([^}]*)\}",
        content,
    )
    if not inp_match:
        return None
    inp_names = [n.strip() for n in inp_match.group(1).split(",") if n.strip()]

    # Find expected pointer array: fp_t const * const exp[N] = { ezz, ed, ... };
    # or: fp_t const * const e[N] = { e0, e1, ... };
    exp_match = re.search(
        r"fp_t\s+const\s*\*\s*const\s+e\w*\[\s*\d*\s*\]\s*=\s*\{([^}]*)\}", content
    )
    if not exp_match:
        return None
    exp_names = [n.strip() for n in exp_match.group(1).split(",") if n.strip()]

    # Build input arrays in order
    inputs = []
    for name in inp_names:
        if name in all_arrays:
            inputs.append(all_arrays[name])
        else:
            return None

    # Build expected arrays in order (may be partial if some can't be resolved)
    expected = []
    for name in exp_names:
        if name in all_arrays:
            expected.append(all_arrays[name])
        # Skip unresolved expected arrays (computed at runtime)

    # Find n
    if inputs:
        n = len(inputs[0])
    else:
        return None

    # Find constants
    constants = []
    # Look for: fp_t threshold = 5; or uint_t k = 3; or similar
    for m in re.finditer(
        r"(?:fp_t|uint_t|int)\s+(?:const\s+)?(\w+)\s*=\s*([^;]+);", content
    ):
        const_name = m.group(1)
        const_val_text = m.group(2).strip()
        # Skip if this is an array declaration
        if "[" in const_val_text:
            continue
        result = _safe_float(const_val_text)
        if result is not None:
            constants.append(result)

    # Find struct parameter: c_t c = { .k = 3, .std_mult = 2.0 };
    params = None
    param_match = re.search(r"(?:c_t|params_t)\s+c\s*=\s*\{([^}]*)\}", content)
    if param_match:
        params = {}
        for fm in re.finditer(r"\.(\w+)\s*=\s*([^,}]+)", param_match.group(1)):
            fname = fm.group(1)
            fval_text = fm.group(2).strip()
            result = _safe_float(fval_text)
            if result is not None:
                params[fname] = result

    # Also look for scalar parameter: uint_t k = 3; fp_t c = 2; fp_t threshold = 0.5;
    if params is None:
        # Try to extract params from ts_test call's 6th argument
        ts_test_match = re.search(
            r"ts_test\s*\([^,]+,\s*[^,]+,\s*[^,]+,\s*[^,]+,\s*[^,]+,\s*(&?(\w+)|NULL)",
            content,
        )
        if ts_test_match:
            param_arg = ts_test_match.group(2)
            if param_arg and param_arg != "NULL":
                # Check if it's an array like "constants[N] = { a, b }"
                arr_match = re.search(
                    rf"fp_t\s+{param_arg}\s*\[\s*\d+\s*\]\s*=\s*\{{([^}}]*)\}}",
                    content,
                )
                if arr_match:
                    var_names = [v.strip() for v in arr_match.group(1).split(",")]
                    resolved = {}
                    for vn in var_names:
                        vn = vn.strip()
                        if not vn:
                            continue
                        result = _safe_float(vn)
                        if result is not None:
                            resolved[vn] = result
                            continue
                        var_match = re.search(
                            rf"(?:uint_t|fp_t|int)\s+(?:const\s+)?{vn}\s*=\s*([^;]+);",
                            content,
                        )
                        if var_match:
                            result = _safe_float(var_match.group(1).strip())
                            if result is not None:
                                resolved[vn] = result
                    if resolved:
                        params = resolved

        if params is None:
            # Extract ALL scalar variables (not just the first)
            scalar_matches = re.finditer(
                r"(?:uint_t|fp_t|int)\s+(?:const\s+)?(\w+)\s*=\s*([^;]+);", content
            )
            collected = {}
            for sm in scalar_matches:
                var_name = sm.group(1)
                if var_name in ("n", "offset", "y_n", "i", "j"):
                    continue
                val_text = sm.group(2).strip()
                result = _safe_float(val_text)
                if result is not None:
                    collected[var_name] = result
            if collected:
                params = collected

    return {
        "inputs": inputs,
        "expected": expected,
        "constants": constants,
        "params": params,
        "n": n,
    }


def call_c_direct(name, x_n, y_n, inputs, n, params=None):
    """Call C flat wrapper directly via ctypes. Returns list of lists or None.

    Parameterized functions (params != None) are skipped — the C test harness
    (make check) validates all 349 functions including those with parameters.
    """
    if params:
        return None
    try:
        flat_func = getattr(ts._lib, f"ts_{name}_flat", None)
        if flat_func is None:
            return None
        np_inputs = [np.array(arr, dtype=np.float64) for arr in inputs]
        y_bufs = [np.empty(n, dtype=np.float64) for _ in range(y_n)]
        args = (
            [b.ctypes.data_as(ctypes.POINTER(ctypes.c_double)) for b in y_bufs]
            + [b.ctypes.data_as(ctypes.POINTER(ctypes.c_double)) for b in np_inputs]
            + [ctypes.c_uint64(n), None]
        )
        ret = flat_func(*args)
        if ret != 0:
            return None
        return [b.tolist() for b in y_bufs]
    except (OSError, AttributeError, ctypes.ArgumentError):
        return None


def isclose(a, b, rel_tol=1e-6, abs_tol=1e-6):
    if math.isnan(a) and math.isnan(b):
        return True
    if math.isinf(a) and math.isinf(b):
        return (a > 0) == (b > 0)
    return abs(a - b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)


def arrays_close(a, b):
    if len(a) != len(b):
        return False
    for x, y in zip(a, b):
        if not isclose(x, y):
            return False
    return True


def normalize_result(result, y_n, n):
    """Normalize Python binding output to list of lists."""
    if isinstance(result, (int, float, np.number)):
        return [[float(result)] * n]
    if isinstance(result, tuple):
        return [r.tolist() if hasattr(r, "tolist") else list(r) for r in result]
    if hasattr(result, "tolist"):
        return [result.tolist()]
    if isinstance(result, list) and len(result) > 0 and hasattr(result[0], "tolist"):
        return [r.tolist() for r in result]
    if isinstance(result, list):
        return result
    return [[float(result)]]


def main():
    status = _run_tests()
    sys.exit(status)


def _run_tests():
    meta = load_func_meta()
    test_dir = os.path.dirname(__file__)

    total = len(meta)
    passed = 0
    failed = 0
    skipped = 0

    sorted_names = sorted(meta.keys())
    rng = __import__("random").Random(SEED)
    # Pre-compute RNG offsets for each function
    rng_offsets = {}
    offset = 0
    for _fname in sorted_names:
        rng_offsets[_fname] = offset
        _x_n = meta[_fname]["x_n"]
        offset += _x_n * N

    for fname, info in sorted(meta.items()):
        x_n = info["x_n"]
        y_n = info["y_n"]

        # Try parsing test file first
        test_path = os.path.join(test_dir, f"{fname}.c")
        parsed = None
        if os.path.exists(test_path):
            parsed = parse_test_file(test_path)

        inputs = None
        expected = None
        params = None
        n = N

        if parsed:
            inputs = parsed["inputs"]
            expected = parsed["expected"]
            params = parsed.get("params")
            n = parsed["n"]

            # Validate input count (expected count is optional - some tests compute at runtime)
            # Note: some C functions declare wrong x_n (e.g., x_n=1 but access x[4]).
            # If test file provides more inputs, trust the test file.
            if len(inputs) < x_n:
                parsed = None
            # Truncate extra inputs to match function's declared x_n
            if len(inputs) > x_n:
                inputs = inputs[:x_n]

        if not parsed:
            # Fallback: random inputs + default params from binding signature
            rng.seed(SEED)
            if hasattr(rng, "jumpahead"):
                rng.jumpahead(rng_offsets.get(fname, 0))
            inputs = [[rng.uniform(0, 100) for _ in range(N)] for _ in range(x_n)]
            n = N
            expected = None
            # Extract default params from the Python binding
            params = get_default_params(fname, x_n)

        # Separate params for C subprocess vs Python binding
        # C subprocess needs original params (may have keys not in binding)
        # Python binding only accepts keys in its signature
        binding_params = get_default_params(fname, x_n)
        if params and binding_params:
            # Filter to keys that exist in binding
            filtered = {k: v for k, v in params.items() if k in binding_params}
            # If no keys match by name, fall back to binding defaults
            c_params = filtered if filtered else binding_params
            py_params = filtered if filtered else binding_params
        elif binding_params:
            # No parsed params but binding has defaults - use them for both
            c_params = binding_params
            py_params = binding_params
        else:
            c_params = params
            py_params = params

        # Call Python binding
        func = getattr(ts, fname, None)
        if func is None:
            skipped += 1
            print(f"  SKIP {fname} (not in Python bindings)")
            continue

        np_inputs = [np.array(arr, dtype=np.float64) for arr in inputs or []]

        # Call C flat wrapper for non-parameterized functions
        actual_x_n = len(inputs) if inputs else x_n
        c_result = call_c_direct(fname, actual_x_n, y_n, inputs, n, c_params)

        try:
            # Pass params to Python binding if available
            if py_params:
                int_params = {}
                for pk, v in py_params.items():
                    int_params[pk] = int(v) if v == int(v) else v
                py_result = func(*np_inputs, **int_params)
            else:
                py_result = func(*np_inputs)
            py_result = normalize_result(py_result, y_n, n)
        except TypeError as e:
            # Param name mismatch: retry with no params (binding defaults)
            if py_params:
                try:
                    py_result = func(*np_inputs)
                    py_result = normalize_result(py_result, y_n, n)
                except Exception as e2:
                    failed += 1
                    print(f"  FAIL {fname}: {e2}")
                    continue
            else:
                failed += 1
                print(f"  FAIL {fname}: {e}")
                continue
        except Exception as e:
            failed += 1
            print(f"  FAIL {fname}: {e}")
            continue

        # If C comparison unavailable (parameterized function), just verify Python didn't crash
        if c_result is None:
            passed += 1
            print(f"  PASS {fname} (py-only)")
            continue

        # Detect C OOB garbage
        _min_normal = sys.float_info.min

        def _is_garbage(arr):
            for v in arr:
                av = abs(v)
                if av == 0:
                    continue
                if 0 < av < _min_normal:
                    return True
            return False

        c_garbage = any(_is_garbage(r) for r in c_result)
        if c_garbage:
            skipped += 1
            print(f"  SKIP {fname} (C OOB garbage)")
            continue

        # Compare against C output
        ok = True
        if len(py_result) != len(c_result):
            ok = False
        else:
            for yi, (got, exp) in enumerate(zip(py_result, c_result)):
                if not arrays_close(got, exp):
                    ok = False
                    for j, (g, e) in enumerate(zip(got, exp)):
                        if not isclose(g, e):
                            print(f"    y[{yi}][{j}]: got={g}, expected={e}")
                            break
                    break

        if ok:
            passed += 1
            print(f"  PASS {fname}")
        else:
            failed += 1
            print(f"  FAIL {fname}")

    # --- Error-case validation (sample) ---
    print("\n--- Error-case validation ---", flush=True)
    error_passed = 0
    error_failed = 0
    error_sample = sorted(meta.keys())[:20]

    # Test 1: Empty arrays
    for fname in error_sample:
        info = meta[fname]
        x_n = info["x_n"]
        func = getattr(ts, fname, None)
        if func is None:
            continue
        try:
            empty_inputs = [np.array([], dtype=np.float64) for _ in range(x_n)]
            sig = inspect.signature(func)
            defaults = {
                k: v.default
                for k, v in sig.parameters.items()
                if v.default is not inspect.Parameter.empty
                and v.annotation is not np.ndarray
            }
            func(*empty_inputs, **defaults) if defaults else func(*empty_inputs)
            error_passed += 1
        except (RuntimeError, ValueError, IndexError, AssertionError):
            error_passed += 1
        except Exception as e:
            error_failed += 1
            print(f"  ERROR {fname} (empty): {type(e).__name__}: {e}")

    # Test 2: Mismatched array lengths
    for fname in error_sample:
        info = meta[fname]
        x_n = info["x_n"]
        if x_n < 2:
            continue
        func = getattr(ts, fname, None)
        if func is None:
            continue
        try:
            mismatched = [
                np.array([1.0] * 10, dtype=np.float64)
                if i == 0
                else np.array([1.0] * 5, dtype=np.float64)
                for i in range(x_n)
            ]
            sig = inspect.signature(func)
            defaults = {
                k: v.default
                for k, v in sig.parameters.items()
                if v.default is not inspect.Parameter.empty
                and v.annotation is not np.ndarray
            }
            func(*mismatched, **defaults) if defaults else func(*mismatched)
            error_failed += 1
            print(f"  ERROR {fname} (mismatched): should have raised")
        except (RuntimeError, ValueError, IndexError, AssertionError):
            error_passed += 1
        except Exception:
            error_passed += 1

    # Test 3: NaN/Inf propagation
    nan_passed = 0
    nan_failed = 0
    for fname in error_sample:
        info = meta[fname]
        x_n = info["x_n"]
        func = getattr(ts, fname, None)
        if func is None:
            continue
        try:
            nan_inputs = [
                np.array([1.0, float("nan"), 3.0, 4.0, 5.0] * 20, dtype=np.float64)
                for _ in range(x_n)
            ]
            sig = inspect.signature(func)
            defaults = {
                k: v.default
                for k, v in sig.parameters.items()
                if v.default is not inspect.Parameter.empty
                and v.annotation is not np.ndarray
            }
            func(*nan_inputs, **defaults) if defaults else func(*nan_inputs)
            nan_passed += 1
        except RuntimeError:
            nan_passed += 1
        except Exception as e:
            nan_failed += 1
            print(f"  ERROR {fname} (NaN): {type(e).__name__}: {e}")

    total_error = error_passed + error_failed
    total_nan = nan_passed + nan_failed
    print(f"\nError-case: {error_passed}/{total_error} passed, {error_failed} failed")
    print(f"NaN/Inf:    {nan_passed}/{total_nan} passed, {nan_failed} failed")

    print(f"\n{passed}/{total} passed, {failed} failed, {skipped} skipped", flush=True)
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    main()