import logging
import uuid
import sys

from itertools import count
from relationalai.early_access.tests.logging import Capturer
from relationalai.early_access.metamodel import ir, types
from relationalai.early_access.builder import builder
from relationalai.clients.result_helpers import sort_data_frame_result
from relationalai.clients.util import IdentityParser
from relationalai.clients.snowflake import Provider as SFProvider
from relationalai import Provider
from typing import cast
from pathlib import Path

class SnapshotTestBase:
    """Base class for tests with snapshots."""

    provider:Provider = cast(SFProvider, Provider()) # type: ignore

    def run_snapshot_test(self, snapshot, script_path, db_schema=None, use_sql=False, use_lqp=False, match_internals=False):
        """Run a snapshot test.

        This method assumes content is already in the correct string format.
        Subclasses should handle any necessary compilation or conversion.
        """
        logger = logging.getLogger("pyrellogger")
        handler = Capturer()
        logger.addHandler(handler)

        # reset state across test runs
        self.reset_state()

        # Track which modules are already loaded
        before_modules = set(sys.modules.keys())
        unique_name = f"md{str(uuid.uuid4())[-12:]}"
        export_fqn = None
        try:
            overrides = {
                'dry_run': False,
                'model_suffix': f"_{unique_name}",
                'use_sql': use_sql,
                'reasoner.rule.use_lqp': use_lqp,
                'keep_model': False,
            }
            with builder.with_overrides(**overrides):
                with open(script_path, "r") as f:
                    # Compile with the correct filename
                    code = compile(f.read(), script_path, "exec")
                    # Execute the compiled code
                    if db_schema:
                        export_fqn = f"{db_schema}.{unique_name}"
                    exec_globals = {"EXPORT_TABLE": export_fqn}
                    exec(code, exec_globals)
        finally:
            # Remove any modules that were imported during exec()
            after_modules = set(sys.modules.keys())
            new_modules = after_modules - before_modules
            for mod in new_modules:
                sys.modules.pop(mod, None)
            # Remove handler
            logger.removeHandler(handler)
            # check if script created an export table
            export_exists = self.is_export_exists(export_fqn)
            # match snapshots
            self.assert_match_snapshots(snapshot, handler, export_exists, export_fqn, use_sql=use_sql, use_lqp=use_lqp, match_internals=match_internals)
            # cleanup resources created during a test run
            self.cleanup(export_exists, export_fqn)

    def assert_match_snapshots(self, snapshot, handler, export_exists=False, export_fqn=None, use_sql=False, use_lqp=False, match_internals=False):
        # compare model and queries only for internal e2e tests
        # todo: remove when e2e merged with unified tests
        if match_internals:
            if use_sql:
                extension = "sql"
            elif use_lqp:
                extension = "lqp"
            else:
                extension = "rel"

            for i, msg in enumerate(handler.model, start=1):
                snapshot.assert_match(str(msg["metamodel"]), f"model_{i}.metamodel")
                snapshot.assert_match(str(msg[extension]), f"model_{i}.{extension}")

            for i, msg in enumerate(handler.queries, start=1):
                snapshot.assert_match(str(msg["metamodel"]), f"query_{i}.metamodel")
                snapshot.assert_match(str(msg[extension]), f"query_{i}.{extension}")

            # The SQL compiler does not use ExtractNestedLogicals, so the type information is for a
            # different IR and we can't really check against what was generated by Rel and LQP
            if not use_sql:
                for i, msg in enumerate(handler.types, start=1):
                    snapshot.assert_match(str(msg["type_report"]), f"type_report_{i}.fish")

        for i, msg in enumerate(handler.results, start=1):
            if "results" in msg:
                snapshot.assert_match(msg["results"].to_csv(index=False).strip('"').strip('\n"'), f"result_{i}.csv")
        if export_fqn and export_exists:
            result_df = sort_data_frame_result(self.provider.sql(f"SELECT * FROM {export_fqn}", format="pandas"))
            snapshot.assert_match(result_df.to_csv(index=False).strip('"').strip('\n"'), "export_result.csv")

    def cleanup(self, export_exists=False, export_fqn=None):
        if export_fqn and export_exists :
            self.provider.sql(f"DROP TABLE IF EXISTS {export_fqn}")

    def is_export_exists(self, export_fqn=None):
        if export_fqn is None:
            return False
        database, schema, table, _ = IdentityParser(export_fqn, require_all_parts=True).to_list()
        return self.provider.sql(f"""
            SELECT EXISTS(
                SELECT 1
                FROM {database}.INFORMATION_SCHEMA.TABLES
                WHERE TABLE_SCHEMA = '{schema}'
                AND TABLE_NAME = '{table}'
            ) AS TABLE_EXISTS;
        """)[0]["TABLE_EXISTS"]

    @staticmethod
    def parent_tests(file: str):
        return SnapshotTestBase.discover_test_scripts(Path(file).parent)

    @staticmethod
    def discover_test_scripts(path: Path):
        scripts_dir = path / "tests"
        return sorted(scripts_dir.glob("**/*.py"))

    @staticmethod
    def get_id(p):
        p_name = p.parent.name
        return p.stem if p_name == "tests" else f"{p_name}_{p.stem}"

    def reset_state(self):
        """
        When we execute a pyrel program we accumulate some state, such as custom decimals, and
        we increase the counter for object ids. This function resets those counters and other
        state to a known baseline to ensure that test snapshots are consistent if we run the
        test alone vs in a suite of tests.
        """
        # reset the global id counters
        ir._global_id = count(10000)
        builder._global_id = count(10000)
        # Used to generate error source IDs
        builder.errors.ModelError.error_locations.clear()
        # caches of custom decimals
        for k in list(types._decimal_types.keys()):
            if types._decimal_types[k] != types.Decimal:
                del types._decimal_types[k]
        for k in list(builder.Concept.builtins):
            if k.startswith("Decimal(") and builder.Concept.builtins[k] != builder.Decimal:
                del builder.Concept.builtins[k]
