"""
Query execution bridge between PyRel frontend and v0 executors.

This module provides a shim layer that translates PyRel frontend queries and models
into v0 metamodel format and executes them using v0 executors (LQP, Rel, DuckDB, or Snowflake).

WARNING: This is an unstable API subject to change.

Key Components:
- ShimExecutor: Main 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
"""

# ------------------------------------------------------
# VERY UNSTABLE API TO EXECUTE FRONTEND PROGRAMS USING V0 EXECUTORS
# ------------------------------------------------------
import json
from functools import lru_cache
from typing import Optional

from pandas import DataFrame

from v0.relationalai import debugging
from v0.relationalai.semantics.lqp.executor import LQPExecutor
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

try:
    from v0.relationalai.clients.snowflake import Provider as v0Provider  # type: ignore
except ImportError:
    from v0.relationalai.clients.resources.snowflake import Provider as v0Provider

from ..config import Config
from ..config.shims import DRY_RUN, FORCE_MODEL_NAME, DATETIME_NOW
from ..semantics import Fragment, Model
from ..semantics.frontend.pprint import format
from ..semantics.metamodel import metamodel as mm
from ..semantics.metamodel.metamodel_analyzer import Normalize
from ..semantics.metamodel.typer import Typer
from .mm2v0 import Translator

PRINT_RESULT = False
TYPER_DEBUGGER = False

# PRINT_RESULT = True
# TYPER_DEBUGGER = True


class ShimExecutor:
    """
    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):
        # if the model is None we derive it from queries
        self.dsl_model = model
        self.normalizer = Normalize()
        # 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
        # the version of the dsl model when the v0_model above was generated
        self.cached_dsl_model_version = -1

    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 V0 e "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 = ShimExecutor(model)
        >>> results = executor.execute(select(Employee.name), executor="lqp")
        """
        if not executor:
            use_lqp = get_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, **source_attrs) as query_span:
            results = self._execute_mm(mm_query, mm_model, executor, export_to=export_to, update=update, meta=meta)
            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_dsl_model_version 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_dsl_model_version = self.dsl_model._version
        self.cached_v0_model = None
        soft_type_errors = get_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 cached_dsl_model_version,
            or if there's no cached model.
        """
        return self.dsl_model and (
            self.cached_dsl_model_version != self.dsl_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):
        """
        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 = self._translate_model(mm_model, debugger_msgs)
            v0_query = self._translate_query(mm_query, 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)

    def _translate_model(self, mm_model: mm.Model | None, debugger_msgs) -> v0.Model | None:
        """
        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.
        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.
        """
        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

        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:
                        mm_model = typer.infer_model(mm_model)
                        if debugging.DEBUG:
                            span["metamodel"] = str(mm_model.root)
                            span["typed_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("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
        return None

    def _translate_query(self, mm_query: mm.Task, 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.
        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:
                        span["metamodel"] = str(mm_query)
                with debugging.span("Typer") as span:
                    mm_query = typer.infer_query(mm_query)
                    if debugging.DEBUG:
                        span["metamodel"] = str(mm_query)
                        span["typed_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:
                        span["metamodel"] = str(mm_query)
                with debugging.span("v0 Translation") as span:
                    v0_query = translator.translate_query(mm_query)
                    if debugging.DEBUG:
                        span["metamodel"] = str(v0_query)
                    assert isinstance(v0_query, v0.Task)
        return v0_query

    def _execute_v0(self, v0_model, v0_query, executor="lqp", export_to="", update=False, meta=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.

        Returns
        -------
        DataFrame
            Query results, or empty DataFrame if in dry-run mode.
        """
        dry_run = DRY_RUN or get_config().compiler.dry_run
        if dry_run and executor not in ("rel", "lqp"):
            results = DataFrame()
        else:
            # create snowflake tables for all the tables that have been used
            assert self.translator is not None
            ts = [v0Table(t.name) for t in self.translator.used_tables if not t.uri.startswith("dataframe://")]
            for t in ts:
                if not any(v0t._fqn == t._fqn for v0t in v0Table._used_sources):
                    t._lazy_init()
                    v0Table._used_sources.add(t)

            export_table = None
            if export_to:
                export_table = v0Table(export_to)

            # get an executor and execute
            database = self.dsl_model.name if self.dsl_model else "placeholder"
            database = FORCE_MODEL_NAME if FORCE_MODEL_NAME else database
            executor = _get_v0_executor(executor, database, dry_run)
            if isinstance(executor, (LQPExecutor, RelExecutor)):
                results = executor.execute(v0_model, v0_query, export_to=export_table, update=update, meta=meta)
            else:
                results = executor.execute(v0_model, v0_query)
            if PRINT_RESULT:
                print(results)
        return results


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}


@lru_cache()
def get_config():
    """
    Get the v0 configuration instance with caching.

    Returns
    -------
    Config
        The v0 configuration object.
    """
    return Config()


@lru_cache()
def _get_v0_executor(name: str, database: str, dry_run: bool):
    """
    Get or create a v0 executor instance with caching.

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

    Parameters
    ----------
    name : str
        Executor type name.
    database : str
        Database name to use.
    dry_run : bool
        Whether to do a dry run or end-2-end.

    Returns
    -------
    Executor
        The appropriate v0 executor instance.

    Raises
    ------
    ValueError
        If executor name is not recognized.
    """
    v0executor.SUPPRESS_TYPE_ERRORS = True
    if name == "duckdb":
        from v0.relationalai.semantics.sql.executor.duck_db import DuckDBExecutor

        return DuckDBExecutor()
    elif name == "lqp":
        # if DATETIME_NOW is set we force a specific datetime for consistent test snapshots
        intrinsic_overrides = {"datetime_now": DATETIME_NOW} if DATETIME_NOW else {}
        return LQPExecutor(database, dry_run=dry_run, intrinsic_overrides=intrinsic_overrides)
    elif name == "rel":
        return RelExecutor(database, dry_run=dry_run)
    elif name == "snowflake":
        from v0.relationalai.semantics.sql.executor.snowflake import SnowflakeExecutor

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


def get_provider():
    """
    Get the v0 Snowflake provider instance.

    Returns
    -------
    v0Provider
        The v0 Snowflake provider for managing Snowflake resources.
    """
    return v0Provider()
