import dataclasses
import difflib
import logging
import typing as t

import patchy.api
import sqlglot
import sqlglot.expressions
import sqlglot.lineage
import sqlglot.optimizer.scope
import sqlglot.optimizer.unnest_subqueries

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 most functions but fails on one:
# - __deepcopy__ in expression_core.py is mypyc-compiled to .so, so we use a wrapper instead

_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_scope_traverse() -> None:
    # Circular scope dependencies can happen in somewhat specific circumstances
    # due to our usage of sqlglot.
    # See https://github.com/tobymao/sqlglot/pull/4244
    #
    # In sqlglot[c] v30+, Scope.traverse is mypyc-compiled and cannot be patched
    # with patchy. We reimplement it entirely in pure Python with cycle detection.
    import itertools

    from sqlglot.errors import OptimizeError
    from sqlglot.optimizer.scope import Scope

    def _traverse_with_cycle_detection(
        self: Scope,
    ) -> t.Iterator[Scope]:
        stack = [self]
        seen_scopes: t.Set[int] = set()
        result = []
        while stack:
            scope = stack.pop()

            # Scopes aren't hashable, so we use id(scope) instead.
            if id(scope) in seen_scopes:
                raise OptimizeError(f"Scope {scope} has a circular scope dependency")
            seen_scopes.add(id(scope))

            result.append(scope)
            stack.extend(
                itertools.chain(
                    scope.cte_scopes,
                    scope.union_scopes,
                    scope.table_scopes,
                    scope.subquery_scopes,
                )
            )

        yield from reversed(result)

    sqlglot.optimizer.scope.Scope.traverse = _traverse_with_cycle_detection  # type: ignore


def _patch_unnest_subqueries() -> None:
    patchy.patch(
        sqlglot.optimizer.unnest_subqueries.decorrelate,
        """\
@@ -142,16 +142,19 @@
         if key in group_by:
             key.replace(nested)
         elif isinstance(predicate, exp.EQ):
-            parent_predicate = _replace(
-                parent_predicate,
-                f"({parent_predicate} AND ARRAY_CONTAINS({nested}, {column}))",
-            )
+            if parent_predicate:
+                parent_predicate = _replace(
+                    parent_predicate,
+                    f"({parent_predicate} AND ARRAY_CONTAINS({nested}, {column}))",
+                )
         else:
             key.replace(exp.to_identifier("_x"))
-            parent_predicate = _replace(
-                parent_predicate,
-                f"({parent_predicate} AND ARRAY_ANY({nested}, _x -> {predicate}))",
-            )
+
+            if parent_predicate:
+                parent_predicate = _replace(
+                    parent_predicate,
+                    f"({parent_predicate} AND ARRAY_ANY({nested}, _x -> {predicate}))",
+                )
""",
    )


def _patch_lineage() -> None:
    # Add the "subfield" attribute to sqlglot.lineage.Node.
    # With dataclasses, the easiest way to do this is with inheritance.
    # 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(sqlglot.lineage.Node):
        subfield: str = ""

    sqlglot.lineage.Node = Node  # type: ignore

    patchy.patch(
        sqlglot.lineage.lineage,
        """\
@@ -95,7 +95,8 @@

     expression = maybe_parse(sql, copy=copy, dialect=dialect)
-    column = normalize_identifiers.normalize_identifiers(column, dialect=dialect).name
+    # column = normalize_identifiers.normalize_identifiers(column, dialect=dialect).name
+    assert isinstance(column, str)

     if sources:
         expression = exp.expand(
""",
    )

    # Patch 1: Change set to list for source_columns
    patchy.patch(
        sqlglot.lineage.to_node,
        """\
@@ -133,9 +133,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: t.Sequence[exp.Expr] = [
""",
    )

    # Patch 2: Add subfield extraction
    patchy.patch(
        sqlglot.lineage.to_node,
        """\
@@ -236,8 +236,21 @@
             # 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()
+
+            subfields = []
+            field: exp.Expression = c
+            while isinstance(field.parent, exp.Dot):
+                field = field.parent
+                subfields.append(field.name)
+            subfield = ".".join(subfields)
+
             node.downstream.append(
-                Node(name=c.sql(comments=False), source=col_expr, expression=col_expr)
+                Node(
+                    name=c.sql(comments=False),
+                    source=col_expr,
+                    expression=col_expr,
+                    subfield=subfield,
+                )
             )

     if _cache is not None:
""",
    )


# Apply patches

# sqlglot.Expression was removed as a top-level re-export in v30; restore it for
# backward compatibility with existing DataHub type annotations.
sqlglot.Expression = sqlglot.expressions.Expression  # type: ignore

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

SQLGLOT_PATCHED = True
