import dataclasses
import difflib
import logging
import typing as t

import patchy.api
import sqlglot
import sqlglot.expressions
import sqlglot.lineage

from datahub.utilities.is_pytest import is_pytest_running
from datahub.utilities.unified_diff import apply_diff

# This injects a few patches into sqlglot to add features and mitigate
# some bugs and performance issues.
# The diffs in this file should match the diffs declared in our fork.
# https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:main
# For a diff-formatted view, see:
# https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:main.diff

# Patches here are generated by running diff -u <original_file> <patched_file>
#
# We use the `patchy` library to apply source-level patches (text diffs) to sqlglot functions.
# With sqlglot[c] (C tokenizer), patchy works for pure-Python functions but fails on those
# compiled by mypyc to .so extensions. For compiled targets we use alternative strategies:
# - __deepcopy__ in expression_core.py: wrapper function
# - lineage.lineage / lineage.to_node: load lineage.py as a pure-Python module and patch that

_DEBUG_PATCHER = is_pytest_running()
logger = logging.getLogger(__name__)

_apply_diff_subprocess = patchy.api._apply_patch


def _new_apply_patch(source: str, patch_text: str, forwards: bool, name: str) -> str:
    assert forwards, "Only forward patches are supported"

    result = apply_diff(source, patch_text)

    # TODO: When in testing mode, still run the subprocess and check that the
    # results line up.
    if _DEBUG_PATCHER:
        result_subprocess = _apply_diff_subprocess(source, patch_text, forwards, name)
        if result_subprocess != result:
            logger.info("Results from subprocess and _apply_diff do not match")
            logger.debug(f"Subprocess result:\n{result_subprocess}")
            logger.debug(f"Our result:\n{result}")
            diff = difflib.unified_diff(
                result_subprocess.splitlines(), result.splitlines()
            )
            logger.debug("Diff:\n" + "\n".join(diff))
            raise ValueError("Results from subprocess and _apply_diff do not match")

    return result


patchy.api._apply_patch = _new_apply_patch

# Store original __deepcopy__ for wrapper
_original_deepcopy = sqlglot.expressions.Expression.__deepcopy__


def _deepcopy_wrapper(
    self: sqlglot.expressions.Expression, memo: t.Any
) -> sqlglot.expressions.Expr:
    """Wrapper for __deepcopy__ with cooperative timeout support.

    Cannot use patchy because expression_core.py is mypyc-compiled to .so in sqlglot[c].
    """
    import datahub.utilities.cooperative_timeout

    datahub.utilities.cooperative_timeout.cooperate()
    return _original_deepcopy(self, memo)


def _patch_lineage() -> None:
    import importlib.util
    import os
    import sys

    # In sqlglot[c] 30.8.0+, lineage.py is compiled by mypyc to a .so extension.
    # patchy cannot patch compiled functions, and mypyc-compiled classes cannot be
    # subclassed by interpreted Python classes. We load the pure Python lineage.py
    # source (installed alongside the .so) and apply our patches to it, then replace
    # the compiled module-level names with our patched Python versions.
    lineage_py_path = os.path.join(
        os.path.dirname(sqlglot.lineage.__file__), "lineage.py"
    )
    spec = importlib.util.spec_from_file_location(
        "sqlglot._lineage_py", lineage_py_path
    )
    assert spec is not None
    lineage_py = importlib.util.module_from_spec(spec)
    # Register in sys.modules before exec so @dataclass can resolve type annotations.
    sys.modules["sqlglot._lineage_py"] = lineage_py
    spec.loader.exec_module(lineage_py)  # type: ignore

    # Add the "subfield" attribute to the Node dataclass.
    # lineage_py.Node is a pure Python dataclass (not compiled), so inheritance works.
    # Unfortunately, mypy won't pick up on the new field, so we need to
    # use type ignores everywhere we use subfield.
    @dataclasses.dataclass(frozen=True)
    class Node(lineage_py.Node):  # type: ignore[name-defined]
        subfield: str = ""

    lineage_py.Node = Node  # type: ignore
    sqlglot.lineage.Node = Node  # type: ignore

    # Patch 1: Change set to list for source_columns (preserve column order)
    patchy.patch(
        lineage_py.to_node,
        """\
@@ -163,9 +163,9 @@

     # Find all columns that went into creating this one to list their lineage nodes.
-    source_columns = set(find_all_in_scope(select, exp.Column))
+    source_columns = list(find_all_in_scope(select, exp.Column))

     # If the source is a UDTF find columns used in the UDTF to generate the table
     if isinstance(source, exp.UDTF):
-        source_columns |= set(source.find_all(exp.Column))
+        source_columns += list(source.find_all(exp.Column))
         derived_tables: Sequence[exp.Expr] = [
""",
    )

    # Patch 2: Add subfield extraction for struct field access (e.g. col.subfield)
    patchy.patch(
        lineage_py.to_node,
        """\
@@ -275,7 +275,20 @@
             # is unknown. This can happen if the definition of a source used in a query is not
             # passed into the `sources` map.
             col_expr = col_source or exp.Placeholder()
-            leaf = Node(name=c.sql(comments=False), source=col_expr, expression=col_expr)
+
+            subfields = []
+            field: exp.Expression = c
+            while isinstance(field.parent, exp.Dot):
+                field = field.parent
+                subfields.append(field.name)
+            subfield = ".".join(subfields)
+
+            leaf = Node(
+                name=c.sql(comments=False),
+                source=col_expr,
+                expression=col_expr,
+                subfield=subfield,
+            )
             node.downstream.append(leaf)
             if on_node:
                 on_node(leaf)
""",
    )

    # Replace the compiled functions with our patched Python versions.
    # lineage_py.lineage calls lineage_py.to_node via its own module namespace,
    # so both patches take effect together.
    sqlglot.lineage.lineage = lineage_py.lineage  # type: ignore
    sqlglot.lineage.to_node = lineage_py.to_node  # type: ignore


# Apply patches

sqlglot.expressions.Expression.__deepcopy__ = _deepcopy_wrapper  # type: ignore
_patch_lineage()

SQLGLOT_PATCHED = True
