"""
Legacy executor: bridges PyRel frontend queries to v0 backend executors.

This module provides the complete backend — type inference, normalization,
v0 translation, and dispatch — for executing PyRel frontend programs.

Key Components:
- LegacyExecutor: Base executor class that handles compilation and execution
- Translation: Converts v1 metamodel to v0 metamodel
- Type inference: Uses Typer to infer types before translation
- Caching: Caches translated models to avoid re-translation

Supported Executors:
- lqp: Logical Query Plan executor
- rel: Rel executor
- duckdb: DuckDB SQL executor
- snowflake: Snowflake SQL executor
"""

from __future__ import annotations
import json
from typing import TYPE_CHECKING, Optional, Union, cast

from pandas import DataFrame

from v0.relationalai import debugging
from v0.relationalai.semantics.lqp.executor import LQPExecutor as V0LQPExecutor, Compiler as LQPCompiler
from v0.relationalai.semantics.metamodel import executor as v0executor
from v0.relationalai.semantics.metamodel import factory as v0_factory
from v0.relationalai.semantics.metamodel import ir as v0
from v0.relationalai.semantics.metamodel.visitor import collect_by_type
from v0.relationalai.semantics.rel.executor import RelExecutor
from v0.relationalai.semantics.snowflake import Table as v0Table
from v0.relationalai.clients.config import Config as V0Config
from v0.relationalai.clients.util import IdentityParser as _IdentityParser

if TYPE_CHECKING:
    from v0.relationalai.semantics.sql.executor.duck_db import DuckDBExecutor
    from v0.relationalai.semantics.sql.executor.snowflake import SnowflakeExecutor


from ...config import create_config
from ...config import shims as _shims
from ...config.v1_to_v0_translator import translate_v1_to_v0
from .. import Fragment, Model
from ..frontend.pprint import format
from ..metamodel import metamodel as mm
from ..metamodel.typer import Typer
from .lqp import rewrite as lqp_rewrite

from ...shims.mm2v0 import Translator

from .executor import Executor as _BaseExecutor
from ..metamodel.metamodel_analyzer import Normalize, Lowering, Groundness

PRINT_RESULT = False
TYPER_DEBUGGER = False

# PRINT_RESULT = True
# TYPER_DEBUGGER = True

# Reserved keyword for (Batch) CSV export - used in calling execute_mm
EXPORT_TO_CSV_RESERVED_TABLE_NAME = "__EXPORT_TO_CSV__"


class _ExportDestination:
    """Lightweight export-only table descriptor.

    Replaces v0Table for export targets so that creating the destination does
    not register CDC sources, Snowflake schema caches, or Rel concepts — all of
    which would cause v0 executors to treat the write target as a read source.
    """
    __slots__ = ("_fqn", "_col_names", "_is_iceberg", "_iceberg_config")

    def __init__(self, fqn: str) -> None:
        _, _, _, self._fqn = _IdentityParser(fqn, require_all_parts=True).to_list()
        self._col_names: list[str] | None = None
        self._is_iceberg: bool = False
        self._iceberg_config = None


class LegacyExecutor(_BaseExecutor):
    """
    Executor that bridges PyRel frontend queries to v0 backend executors.

    This class handles the full compilation and execution pipeline:
    1. Converts frontend Fragment/Model to v1 metamodel
    2. Type inference using Typer
    3. Normalization
    4. Translation to v0 metamodel
    5. Execution using selected v0 executor

    The executor maintains caches to avoid re-translating unchanged models.

    Attributes
    ----------
    dsl_model : Model | None
        The PyRel model to execute queries against. If None, a v0 model is derived from
        queries.
    """

    def __init__(self, model: Model | None = None):
        super().__init__()
        # if the model is None we derive it from queries
        self.dsl_model = model
        self._normalizer = Normalize()
        self._groundness = Groundness()
        # typer and translator are created on demand when the model changes
        self._typer = None
        self.translator = None
        # the last v0 model generated for the dsl_model
        self.cached_v0_model = None
        # (id(model), model._version) when the v0 model above was generated
        self._cached_prepared_model_version = None
        self._cached_prepared_mm_model = None
        self.config = model.config if model else create_config()
        self.v0_config = V0Config(translate_v1_to_v0(self.config), fetch=False)
        self._executor_cache = {}

    def execute_query(self, model, fragment, **kwargs) -> DataFrame:
        return self.execute(fragment, export_to=kwargs.get("export_to", ""), update=kwargs.get("update", False), meta=kwargs.get("meta"))

    def execute(self, query: Fragment, executor: Optional[str] = None, export_to="", update=False, meta=None):
        """
        Execute a PyRel frontend query.

        This is the main entry point for executing queries. It converts the frontend
        Fragment to metamodel format, then to v0, and then executes it using the specified
        executor.

        Parameters
        ----------
        query : Fragment
            The PyRel query fragment to execute.
        executor : str | None
            V0 executor: "lqp", "rel", "duckdb", or "snowflake". If None, uses
            the configured default (lqp if use_lqp is true, else rel).
        export_to : str
            Optional table name to export results to.
        update : bool
            Whether this is an update operation.
        meta : Any
            Optional metadata to pass to the executor.

        Returns
        -------
        DataFrame
            Query results as a pandas DataFrame.

        Examples
        --------
        Execute a query using LQP executor:

        >>> executor = LegacyExecutor(model)
        >>> results = executor.execute(select(Employee.name), executor="lqp")
        """
        if not executor:
            use_lqp = self.config.reasoners.logic.use_lqp
            executor = "lqp" if use_lqp else "rel"

        if self._should_translate_model():
            assert isinstance(self.dsl_model, Model)
            mm_model = self.dsl_model.to_metamodel()
        else:
            mm_model = None
        mm_query = query.to_metamodel()
        assert isinstance(mm_query, mm.Task)
        return self.execute_mm(
            format(query), mm_query, mm_model, executor, export_to=export_to, update=update, meta=meta
        )

    def execute_mm(
        self,
        dsl_query_string: str,
        mm_query: mm.Task,
        mm_model: mm.Model | None = None,
        executor="lqp",
        export_to="",
        update=False,
        meta=None,
    ):
        """
        Execute a query that's already in v1 metamodel format.

        This method wraps the execution in a debugging span and includes source
        information for debugging.

        Parameters
        ----------
        dsl_query_string : str
            String representation of the query for debugging.
        mm_query : mm.Task
            The query in v1 metamodel format.
        mm_model : mm.Model | None
            Optional model in v1 metamodel format. If None, uses cached model.
        executor : str
            V0 executor to use: "lqp", "rel", "duckdb", or "snowflake".
        export_to : str
            Optional table name to export results to.
        update : bool
            Whether this is an update operation.
        meta : Any
            Optional metadata to pass to the executor.

        Returns
        -------
        DataFrame
            Query results as a pandas DataFrame.
        """
        source_attrs = with_source(mm_query)
        # If there's no source file, use the query_string as the source so it shows up in the debugger
        if not source_attrs.get("source"):
            source_attrs["source"] = dsl_query_string

        with debugging.span("query", tag=None, export_to=export_to, dsl=dsl_query_string, meta=meta, **source_attrs) as query_span:
            results = self._execute_mm(mm_query, mm_model, executor, export_to=export_to, update=update, meta=meta, source_info=source_attrs)
            query_span["results"] = results
            return results

    # ------------------------------------------------------
    # Internal
    # ------------------------------------------------------

    def _reset(self) -> tuple[Typer, Translator]:
        """
        Reset internal state with fresh typer and translator.

        Called when the model changes to ensure clean compilation state.
        Updates the cached version token to track the new model version.

        Returns
        -------
        tuple[Typer, Translator]
            Newly created typer and translator instances.
        """
        if self.dsl_model is not None:
            self._cached_prepared_model_version = (id(self.dsl_model), self.dsl_model._version)
        self.cached_v0_model = None
        self._cached_prepared_mm_model = None
        soft_type_errors = self.config.compiler.soft_type_errors
        # Note: 'enforce' in Typer means strict mode (fail on errors), opposite of soft_type_errors
        self._typer = Typer(enforce=not soft_type_errors)
        self.translator = Translator()
        return self._typer, self.translator

    def _should_translate_model(self):
        """
        Check if the DSL model has changed since last compilation or if the cached v0 model
        is missing (e.g. due to exceptions in the previous translation).

        Returns
        -------
        bool
            True if there is a dsl_model and its version differs from the cached version token,
            or if there's no cached model.
        """
        if not self.dsl_model:
            return False
        current_token = (id(self.dsl_model), self.dsl_model._version)
        return current_token != self._cached_prepared_model_version or self.cached_v0_model is None

    def _execute_mm(self, mm_query: mm.Task, mm_model: mm.Model | None, executor, export_to, update, meta, source_info):
        """
        Execute a query in v1 metamodel format.

        Translates both model and query to v0 format, then executes using
        the specified v0 executor. If mm_model is None and no cached model
        exists, derives a minimal model from the query.

        Parameters
        ----------
        mm_query : mm.Task
            The query in v1 metamodel format.
        mm_model : mm.Model | None
            Optional model in v1 metamodel format.
        executor : str
            V0 executor to use.
        export_to : str
            Optional table name for export.
        update : bool
            Whether this is an update operation.
        meta : Any
            Optional metadata for the executor.

        Returns
        -------
        DataFrame
            Query results.
        """

        # for typer debugging
        debugger_msgs = []
        try:
            v0_model, new_v0_model = self._translate_model(mm_model, executor, debugger_msgs)
            v0_query = self._translate_query(mm_query, executor, debugger_msgs)
            if v0_model is None:
                # there was no model, so create one from the elements refered to by the query
                v0_model = v0_factory.model(
                    collect_by_type(v0.Engine, v0_query),
                    collect_by_type(v0.Relation, v0_query),
                    v0_factory._collect_reachable_types(collect_by_type(v0.Type, v0_query)),
                    v0_factory.logical([]),
                )
        finally:
            if TYPER_DEBUGGER:
                with open("typer_debug.jsonl", "w") as f:
                    for msg in debugger_msgs:
                        f.write(msg)
                        f.write("\n")

        return self._execute_v0(v0_model, v0_query, executor=executor, export_to=export_to, update=update, meta=meta, new_v0_model=new_v0_model, source_info=source_info)

    def _translate_model(self, mm_model: mm.Model | None, executor, debugger_msgs) -> tuple[v0.Model | None, bool]:
        """
        Translate a v1 metamodel model to v0 format with caching.

        The translation pipeline includes:
        1. Type inference using Typer
        2. Normalization
        3. Translation to v0 metamodel

        If mm_model is None, returns the cached v0 model if available.

        Parameters
        ----------
        mm_model : mm.Model | None
            The v1 metamodel model to translate. If None, indicates model hasn't
            changed and cached translation should be reused.
        executor : str
            V0 executor to use.
        debugger_msgs : list
            List to append debugging messages to for the typer debugger.

        Returns
        -------
        v0.Model | None
            The translated v0 model, cached v0 model, or None if no model available.
        bool
            True if a new v0 model was generated, False if cached model was used.
        """
        if mm_model is None and self.cached_v0_model is not None:
            # if the model has not changed we use the previous v0 model
            return self.cached_v0_model, False

        elif mm_model is not None:
            # model has changed since last compilation
            typer, translator = self._reset()
            debugger_msgs.append(json.dumps({"id": "model", "content": str(mm_model)}))
            with debugging.span("compile", metamodel=mm_model) as span:
                span["compile_type"] = "model_v1"
                with debugging.span("v1") as span:
                    with debugging.span("Original") as span:
                        if debugging.DEBUG:
                            span["metamodel"] = str(mm_model.root)
                    with debugging.span("Typer") as span:
                        try:
                            mm_model = typer.infer_model(mm_model)
                        finally:
                            if debugging.DEBUG:
                                span["metamodel"] = str(mm_model.root)
                                span["typed_model_mm"] = str(mm_model)
                                debugger_msgs.append(json.dumps({"id": "typed_model", "content": str(mm_model)}))
                                assert typer.model_net is not None
                                debugger_msgs.append(
                                    json.dumps(
                                        {
                                            "id": "model_net",
                                            "content": str(typer.model_net.to_mermaid()),
                                        }
                                    )
                                )
                    with debugging.span("Normalizer") as span:
                        mm_model = mm_model.mut(root=self._normalizer.normalize(mm_model.root))  # type: ignore
                        assert isinstance(mm_model, mm.Model)
                        if debugging.DEBUG:
                            span["metamodel"] = str(mm_model.root)
                    with debugging.span("Lowering") as span:
                        mm_model = mm_model.mut(root=Lowering().normalize(mm_model.root))  # type: ignore
                        assert isinstance(mm_model, mm.Model)
                        if debugging.DEBUG:
                            span["metamodel"] = str(mm_model.root)
                    with debugging.span("Groundness") as span:
                        self._groundness.enforce(mm_model.root)
                    mm_model = self._rewrite_model(mm_model)
                with debugging.span("v1_v0") as span:
                    with debugging.span("v0 Translation") as span:
                        self.cached_v0_model = translator.translate_model(mm_model)
                        if debugging.DEBUG:
                            span["metamodel"] = str(self.cached_v0_model.root)
                        return self.cached_v0_model, True
        return None, True

    def _translate_query(self, mm_query: mm.Task, executor, debugger_msgs) -> v0.Task:
        """
        Translate a v1 metamodel query to v0 format.

        The translation pipeline includes:
        1. Type inference using Typer
        2. Normalization
        3. Translation to v0 metamodel

        Parameters
        ----------
        mm_query : mm.Task
            The v1 metamodel query to translate.
        executor : str
            The v0 executor to use.
        debugger_msgs : list
            List to append debugging messages to for the typer debugger.

        Returns
        -------
        v0.Task
            The translated v0 query.
        """
        if self._typer is None or self.translator is None:
            typer, translator = self._reset()
        else:
            typer, translator = self._typer, self.translator

        debugger_msgs.append(json.dumps({"id": "query", "content": str(mm_query)}))
        with debugging.span("compile", metamodel=mm_query) as span:
            span["compile_type"] = "query_v1"
            with debugging.span("v1") as span:
                with debugging.span("Original") as span:
                    if debugging.DEBUG or TYPER_DEBUGGER:
                        span["metamodel"] = str(mm_query)
                with debugging.span("Typer") as span:
                    try:
                        mm_query = typer.infer_query(mm_query)
                    finally:
                        if debugging.DEBUG or TYPER_DEBUGGER:
                            span["metamodel"] = str(mm_query)
                            span["typed_query_mm"] = str(mm_query)
                            debugger_msgs.append(json.dumps({"id": "typed_query", "content": str(mm_query)}))
                            assert typer.last_net is not None
                            debugger_msgs.append(
                                json.dumps(
                                    {
                                        "id": "query_net",
                                        "content": str(typer.last_net.to_mermaid()),
                                    }
                                )
                            )
                with debugging.span("Normalizer") as span:
                    mm_query = self._normalizer.normalize(mm_query)  # type: ignore
                    assert isinstance(mm_query, mm.Task)
                    if debugging.DEBUG or TYPER_DEBUGGER:
                        span["metamodel"] = str(mm_query)
                with debugging.span("Lowering") as span:
                    mm_query = Lowering().normalize(mm_query)  # type: ignore
                    assert isinstance(mm_query, mm.Task)
                    if debugging.DEBUG or TYPER_DEBUGGER:
                        span["metamodel"] = str(mm_query)
                with debugging.span("Groundness") as span:
                    self._groundness.enforce(mm_query)
                mm_query = self._rewrite_query(mm_query)
            with debugging.span("v1_v0") as span:
                with debugging.span("v0 Translation") as span:
                    v0_query = translator.translate_query(mm_query)
                    if debugging.DEBUG or TYPER_DEBUGGER:
                        span["metamodel"] = str(v0_query)
                    assert isinstance(v0_query, v0.Task)
        return v0_query

    def _rewrite_model(self, mm_model: mm.Model) -> mm.Model:
        """Post-normalize model rewrite hook. Override in subclasses to apply passes."""
        return mm_model

    def _rewrite_query(self, mm_query: mm.Task) -> mm.Task:
        """Post-normalize query rewrite hook. Override in subclasses to apply passes."""
        return mm_query

    def _execute_v0(self, v0_model, v0_query, executor="lqp", export_to="", update=False, meta=None, new_v0_model=False, source_info=None):
        """
        Execute a query in v0 metamodel format using a v0 executor.

        Handles initialization of Snowflake tables if needed and delegates to
        the appropriate v0 executor.

        Parameters
        ----------
        v0_model : v0.Model
            The model in v0 metamodel format.
        v0_query : v0.Task
            The query in v0 metamodel format.
        executor : str
            Executor name: "lqp", "rel", "duckdb", or "snowflake".
        export_to : str
            Optional table name for export.
        update : bool
            Whether this is an update operation.
        meta : Any
            Optional metadata for the executor.
        new_v0_model : bool
            Whether a new v0 model was generated for this execution, in which case we should
            discard any caches.

        Returns
        -------
        DataFrame
            Query results, or empty DataFrame if in dry-run mode.
        """
        dry_run = _shims.DRY_RUN or self.config.compiler.dry_run
        if dry_run and executor not in ("rel", "lqp"):
            results = DataFrame()
        else:
            session = self.config.get_session() if not dry_run else None

            # create snowflake tables for all the tables that have been used
            assert self.translator is not None

            for source_table in self.translator.used_tables:
                if source_table.uri.startswith("dataframe://"):
                    continue
                # Skip the export destination — it is a write target, not a CDC source.
                # The export destination may appear in used_tables because its columns are
                # referenced in the query for output mapping. Treating it as a source
                # would cause _lazy_init() to run SHOW COLUMNS on a table that may not
                # yet exist
                if (export_to
                        and export_to != EXPORT_TO_CSV_RESERVED_TABLE_NAME
                        and source_table.name == export_to):
                    continue

                if not session:
                    raise RuntimeError("Tables are not supported in dry_run mode.")
                else:
                    from snowflake.snowpark.session import Session as SnowparkSession

                    if not isinstance(session, SnowparkSession):
                        raise RuntimeError("Tables are supported only in a snowflake connection.")

                table = v0Table(source_table.name)
                table._skip_cdc = source_table.skip_cdc

                if any(v0t._fqn == table._fqn for v0t in v0Table._used_sources):
                    continue

                table._lazy_init()

                # When `skip_cdc` is set, the underlying table is
                # expected to be loaded into the GI database
                # already.
                if source_table.skip_cdc:
                    continue

                v0Table._used_sources.add(table)

            export_table = None
            if export_to and export_to != EXPORT_TO_CSV_RESERVED_TABLE_NAME:
                export_table = _ExportDestination(export_to)

            # get an executor and execute
            database = self.dsl_model.name if self.dsl_model else "placeholder"
            database = _shims.FORCE_MODEL_NAME if _shims.FORCE_MODEL_NAME else database
            v0_executor = self._get_v0_executor(executor, database, dry_run)
            if isinstance(v0_executor, V0LQPExecutor):
                if new_v0_model:
                    # reset the LQP compiler because it can be cached and contain stale data for the model
                    v0_executor.compiler = LQPCompiler(passes=lqp_rewrite.v0_lqp_passes())
                format = "csv" if isinstance(export_to, str) and export_to == EXPORT_TO_CSV_RESERVED_TABLE_NAME else "pandas"
                results = v0_executor.execute(v0_model, v0_query, format, cast(v0Table, export_table), update, meta, source_info=source_info)
            elif isinstance(v0_executor, RelExecutor):
                results = v0_executor.execute(v0_model, v0_query, export_to=cast(v0Table, export_table), update=update, meta=meta, source_info=source_info)
            else:
                results = v0_executor.execute(v0_model, v0_query, source_info=source_info)
            if PRINT_RESULT:
                print(results)
        return results

    def _get_v0_executor(self, name: str, database: str, dry_run: bool) -> Union[V0LQPExecutor, RelExecutor, "DuckDBExecutor", "SnowflakeExecutor"]:
        cache_key = (name, database, dry_run)
        if cache_key in self._executor_cache:
            return self._executor_cache[cache_key]

        v0executor.SUPPRESS_TYPE_ERRORS = True  # type: ignore[attr-defined]
        v0_config = self.v0_config

        if name == "duckdb":
            from v0.relationalai.semantics.sql.executor.duck_db import DuckDBExecutor

            executor = DuckDBExecutor()
        elif name == "lqp":
            # if DATETIME_NOW is set we force a specific datetime for consistent test snapshots
            intrinsic_overrides = {"datetime_now": _shims.DATETIME_NOW} if _shims.DATETIME_NOW else {}
            executor = V0LQPExecutor(database, dry_run=dry_run, config=v0_config, intrinsic_overrides=intrinsic_overrides)
            executor.compiler = LQPCompiler(passes=lqp_rewrite.v0_lqp_passes())
        elif name == "rel":
            executor = RelExecutor(database, dry_run=dry_run, config=v0_config)
        elif name == "snowflake":
            from v0.relationalai.semantics.sql.executor.snowflake import SnowflakeExecutor

            executor = SnowflakeExecutor(database="TEST_DB", schema="PUBLIC")
        else:
            raise ValueError(f"Unknown executor: {name}")

        self._executor_cache[cache_key] = executor
        return executor


def with_source(item: mm.Node):
    """
    Extract source location information from a metamodel node.

    Used for debugging to provide file, line number, and source code context.

    Parameters
    ----------
    item : mm.Node
        The metamodel node to extract source information from.

    Returns
    -------
    dict
        Dictionary with source information keys: "file", "line", optionally "source".
        Empty dict if source is None.

    Raises
    ------
    ValueError
        If the item has no source attribute.
    """
    if not hasattr(item, "source"):
        raise ValueError(f"Item {item} has no source")
    elif item.source is None:
        return {}
    elif debugging.DEBUG:
        source = item.source.block
        if source:
            return {"file": source.file, "line": source.line, "source": source.source}
        else:
            return {"file": item.source.file, "line": item.source.line}
    else:
        return {"file": item.source.file, "line": item.source.line}
