import sys
from collections.abc import Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Optional, Union, cast

if TYPE_CHECKING:
    from alembic.migration import MigrationContext
    from alembic.operations.ops import MigrationScript, UpgradeOps
    from click import Group

    from advanced_alchemy.config import SQLAlchemyAsyncConfig, SQLAlchemySyncConfig

__all__ = ("add_migration_commands", "get_alchemy_group")


def get_alchemy_group() -> "Group":
    """Get the Advanced Alchemy CLI group.

    Raises:
        MissingDependencyError: If the `click` package is not installed.

    Returns:
        The Advanced Alchemy CLI group.
    """
    from advanced_alchemy.exceptions import MissingDependencyError

    try:
        import rich_click as click
    except ImportError:
        try:
            import click  # type: ignore[no-redef]
        except ImportError as e:
            raise MissingDependencyError(package="click", install_package="cli") from e

    @click.group(name="alchemy")
    @click.option(
        "--config",
        help="Dotted path to SQLAlchemy config(s) (e.g. 'myapp.config.alchemy_configs')",
        required=True,
        type=str,
    )
    @click.pass_context
    def alchemy_group(ctx: "click.Context", config: str) -> None:
        """Advanced Alchemy CLI commands."""
        from pathlib import Path

        from rich import get_console

        from advanced_alchemy.utils import module_loader

        console = get_console()
        ctx.ensure_object(dict)

        # Add current working directory to sys.path to allow loading local config modules
        cwd = str(Path.cwd())
        if cwd not in sys.path:
            sys.path.insert(0, cwd)

        try:
            config_instance = module_loader.import_string(config)
            if isinstance(config_instance, Sequence):
                ctx.obj["configs"] = config_instance
            else:
                ctx.obj["configs"] = [config_instance]
        except ImportError as e:
            console.print(f"[red]Error loading config: {e}[/]")
            ctx.exit(1)
        finally:
            # Clean up: remove the cwd from sys.path if we added it
            if cwd in sys.path and sys.path[0] == cwd:
                sys.path.remove(cwd)

    return alchemy_group


def add_migration_commands(database_group: Optional["Group"] = None) -> "Group":  # noqa: C901, PLR0915
    """Add migration commands to the database group.

    Args:
        database_group: The database group to add the commands to.

    Raises:
        MissingDependencyError: If the `click` package is not installed.

    Returns:
        The database group with the migration commands added.
    """
    from advanced_alchemy.exceptions import MissingDependencyError

    try:
        import rich_click as click
    except ImportError:
        try:
            import click  # type: ignore[no-redef]
        except ImportError as e:
            raise MissingDependencyError(package="click", install_package="cli") from e
    from rich import get_console

    console = get_console()

    if database_group is None:
        database_group = get_alchemy_group()

    bind_key_option = click.option(
        "--bind-key",
        help="Specify which SQLAlchemy config to use by bind key",
        type=str,
        default=None,
    )
    verbose_option = click.option(
        "--verbose",
        help="Enable verbose output.",
        type=bool,
        default=False,
        is_flag=True,
    )
    no_prompt_option = click.option(
        "--no-prompt",
        help="Do not prompt for confirmation before executing the command.",
        type=bool,
        default=False,
        required=False,
        show_default=True,
        is_flag=True,
    )

    def get_config_by_bind_key(
        ctx: "click.Context", bind_key: Optional[str]
    ) -> "Union[SQLAlchemyAsyncConfig, SQLAlchemySyncConfig]":
        """Get the SQLAlchemy config for the specified bind key.

        Args:
            ctx: The click context.
            bind_key: The bind key to get the config for.

        Returns:
            The SQLAlchemy config for the specified bind key.
        """
        configs = ctx.obj["configs"]
        if bind_key is None:
            return cast("Union[SQLAlchemyAsyncConfig, SQLAlchemySyncConfig]", configs[0])

        for config in configs:
            if config.bind_key == bind_key:
                return cast("Union[SQLAlchemyAsyncConfig, SQLAlchemySyncConfig]", config)

        console.print(f"[red]No config found for bind key: {bind_key}[/]")
        sys.exit(1)

    @database_group.command(
        name="show-current-revision",
        help="Shows the current revision for the database.",
    )
    @bind_key_option
    @verbose_option
    def show_database_revision(bind_key: Optional[str], verbose: bool) -> None:  # pyright: ignore[reportUnusedFunction]
        """Show current database revision."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Listing current revision[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.current(verbose=verbose)

    @database_group.command(
        name="downgrade",
        help="Downgrade database to a specific revision.",
    )
    @bind_key_option
    @click.option("--sql", type=bool, help="Generate SQL output for offline migrations.", default=False, is_flag=True)
    @click.option(
        "--tag",
        help="an arbitrary 'tag' that can be intercepted by custom env.py scripts via the .EnvironmentContext.get_tag_argument method.",
        type=str,
        default=None,
    )
    @no_prompt_option
    @click.argument(
        "revision",
        type=str,
        default="-1",
    )
    def downgrade_database(  # pyright: ignore[reportUnusedFunction]
        bind_key: Optional[str], revision: str, sql: bool, tag: Optional[str], no_prompt: bool
    ) -> None:
        """Downgrade the database to the latest revision."""
        from rich.prompt import Confirm

        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Starting database downgrade process[/]", align="left")
        input_confirmed = (
            True
            if no_prompt
            else Confirm.ask(f"Are you sure you want to downgrade the database to the `{revision}` revision?")
        )
        if input_confirmed:
            sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
            alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
            alembic_commands.downgrade(revision=revision, sql=sql, tag=tag)

    @database_group.command(
        name="upgrade",
        help="Upgrade database to a specific revision.",
    )
    @bind_key_option
    @click.option("--sql", type=bool, help="Generate SQL output for offline migrations.", default=False, is_flag=True)
    @click.option(
        "--tag",
        help="an arbitrary 'tag' that can be intercepted by custom env.py scripts via the .EnvironmentContext.get_tag_argument method.",
        type=str,
        default=None,
    )
    @no_prompt_option
    @click.argument(
        "revision",
        type=str,
        default="head",
    )
    def upgrade_database(  # pyright: ignore[reportUnusedFunction]
        bind_key: Optional[str], revision: str, sql: bool, tag: Optional[str], no_prompt: bool
    ) -> None:
        """Upgrade the database to the latest revision."""
        from rich.prompt import Confirm

        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Starting database upgrade process[/]", align="left")
        input_confirmed = (
            True
            if no_prompt
            else Confirm.ask(f"[bold]Are you sure you want migrate the database to the `{revision}` revision?[/]")
        )
        if input_confirmed:
            sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
            alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
            alembic_commands.upgrade(revision=revision, sql=sql, tag=tag)

    @database_group.command(
        help="Stamp the revision table with the given revision; don't run any migrations",
    )
    @click.argument("revision", type=str)
    @bind_key_option
    @click.option("--sql", is_flag=True, default=False, help="Generate SQL output for offline migrations")
    @click.option(
        "--tag", type=str, default=None, help="Arbitrary 'tag' that can be intercepted by custom env.py scripts"
    )
    @click.option("--purge", is_flag=True, default=False, help="Delete all entries in version table before stamping")
    def stamp(bind_key: Optional[str], revision: str, sql: bool, tag: Optional[str], purge: bool) -> None:  # pyright: ignore[reportUnusedFunction]
        """Stamp the revision table with the given revision."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Stamping revision table[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.stamp(revision=revision, sql=sql, tag=tag, purge=purge)

    @database_group.command(
        name="check",
        help="Check if the target database is up to date",
    )
    @bind_key_option
    def check_revision(bind_key: Optional[str]) -> None:  # pyright: ignore[reportUnusedFunction]
        """Check for pending upgrade operations."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Checking for pending migrations[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.check()

    @database_group.command(
        name="edit",
        help="Edit a revision file using $EDITOR",
    )
    @click.argument("revision", type=str)
    @bind_key_option
    def edit_revision(bind_key: Optional[str], revision: str) -> None:  # pyright: ignore[reportUnusedFunction]
        """Edit revision script with system editor."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule(f"[yellow]Opening revision {revision} in editor[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.edit(revision=revision)

    @database_group.command(
        name="ensure-version",
        help="Create the alembic version table if it doesn't exist",
    )
    @bind_key_option
    @click.option("--sql", is_flag=True, default=False, help="Generate SQL output instead of executing")
    def ensure_version_table(bind_key: Optional[str], sql: bool) -> None:  # pyright: ignore[reportUnusedFunction]
        """Ensure alembic version table exists."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Ensuring version table exists[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.ensure_version(sql=sql)

    @database_group.command(
        name="heads",
        help="Show current available heads in the script directory",
    )
    @bind_key_option
    @verbose_option
    @click.option(
        "--resolve-dependencies",
        is_flag=True,
        default=False,
        help="Resolve dependencies between heads",
    )
    def show_heads(bind_key: Optional[str], verbose: bool, resolve_dependencies: bool) -> None:  # pyright: ignore[reportUnusedFunction]
        """Show current heads."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Showing current heads[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.heads(verbose=verbose, resolve_dependencies=resolve_dependencies)

    @database_group.command(
        name="history",
        help="List changeset scripts in chronological order",
    )
    @bind_key_option
    @verbose_option
    @click.option(
        "--rev-range",
        type=str,
        default=None,
        help="Revision range (e.g., 'base:head', 'abc:def')",
    )
    @click.option(
        "--indicate-current",
        is_flag=True,
        default=False,
        help="Indicate the current revision",
    )
    def show_history(  # pyright: ignore[reportUnusedFunction]
        bind_key: Optional[str],
        verbose: bool,
        rev_range: Optional[str],
        indicate_current: bool,
    ) -> None:
        """Show revision history."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Showing revision history[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.history(
            rev_range=rev_range,
            verbose=verbose,
            indicate_current=indicate_current,
        )

    @database_group.command(
        name="merge",
        help="Merge two revisions together, creating a new migration file",
    )
    @click.argument("revisions", type=str)
    @bind_key_option
    @click.option("-m", "--message", type=str, default=None, help="Merge message")
    @click.option("--branch-label", type=str, default=None, help="Branch label for merge revision")
    @click.option("--rev-id", type=str, default=None, help="Specify custom revision ID")
    @no_prompt_option
    def merge_revisions(  # pyright: ignore[reportUnusedFunction]
        bind_key: Optional[str],
        revisions: str,
        message: Optional[str],
        branch_label: Optional[str],
        rev_id: Optional[str],
        no_prompt: bool,
    ) -> None:
        """Merge revisions (resolves multiple heads)."""
        from rich.prompt import Prompt

        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Merging revisions[/]", align="left")

        if message is None:
            message = "merge revisions" if no_prompt else Prompt.ask("Enter merge message")

        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.merge(
            revisions=revisions,
            message=message,
            branch_label=branch_label,
            rev_id=rev_id,
        )

    @database_group.command(
        name="show",
        help="Show the revision denoted by the given symbol",
    )
    @click.argument("revision", type=str)
    @bind_key_option
    def show_revision(bind_key: Optional[str], revision: str) -> None:  # pyright: ignore[reportUnusedFunction]
        """Show details of a specific revision."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule(f"[yellow]Showing revision {revision}[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.show(rev=revision)

    @database_group.command(
        name="branches",
        help="Show current branch points in the migration history",
    )
    @bind_key_option
    @verbose_option
    def show_branches(bind_key: Optional[str], verbose: bool) -> None:  # pyright: ignore[reportUnusedFunction]
        """Show branch points."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Showing branch points[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.branches(verbose=verbose)

    @database_group.command(
        name="list-templates",
        help="List available Alembic migration templates",
    )
    @bind_key_option
    def list_init_templates(bind_key: Optional[str]) -> None:  # pyright: ignore[reportUnusedFunction]
        """List available initialization templates."""
        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Available templates[/]", align="left")
        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.list_templates()

    @database_group.command(
        name="init",
        help="Initialize migrations for the project.",
    )
    @bind_key_option
    @click.argument(
        "directory",
        default=None,
        required=False,
    )
    @click.option("--multidb", is_flag=True, default=False, help="Support multiple databases")
    @click.option("--package", is_flag=True, default=True, help="Create `__init__.py` for created folder")
    @no_prompt_option
    def init_alembic(  # pyright: ignore[reportUnusedFunction]
        bind_key: Optional[str], directory: Optional[str], multidb: bool, package: bool, no_prompt: bool
    ) -> None:
        """Initialize the database migrations."""
        from rich.prompt import Confirm

        from advanced_alchemy.alembic.commands import AlembicCommands

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Initializing database migrations.", align="left")
        input_confirmed = (
            True if no_prompt else Confirm.ask("[bold]Are you sure you want initialize migrations for the project?[/]")
        )
        if input_confirmed:
            configs = [get_config_by_bind_key(ctx, bind_key)] if bind_key is not None else ctx.obj["configs"]
            for config in configs:
                directory = config.alembic_config.script_location if directory is None else directory
                alembic_commands = AlembicCommands(sqlalchemy_config=config)
                alembic_commands.init(directory=cast("str", directory), multidb=multidb, package=package)

    @database_group.command(
        name="make-migrations",
        help="Create a new migration revision.",
    )
    @bind_key_option
    @click.option("-m", "--message", default=None, help="Revision message")
    @click.option(
        "--autogenerate/--no-autogenerate", default=True, help="Automatically populate revision with detected changes"
    )
    @click.option("--sql", is_flag=True, default=False, help="Export to `.sql` instead of writing to the database.")
    @click.option("--head", default="head", help="Specify head revision to use as base for new revision.")
    @click.option(
        "--splice", is_flag=True, default=False, help='Allow a non-head revision as the "head" to splice onto'
    )
    @click.option("--branch-label", default=None, help="Specify a branch label to apply to the new revision")
    @click.option("--version-path", default=None, help="Specify specific path from config for version file")
    @click.option("--rev-id", default=None, help="Specify a ID to use for revision.")
    @no_prompt_option
    def create_revision(  # pyright: ignore[reportUnusedFunction]
        bind_key: Optional[str],
        message: Optional[str],
        autogenerate: bool,
        sql: bool,
        head: str,
        splice: bool,
        branch_label: Optional[str],
        version_path: Optional[str],
        rev_id: Optional[str],
        no_prompt: bool,
    ) -> None:
        """Create a new database revision."""
        from rich.prompt import Prompt

        from advanced_alchemy.alembic.commands import AlembicCommands

        def process_revision_directives(
            context: "MigrationContext",  # noqa: ARG001
            revision: tuple[str],  # noqa: ARG001
            directives: list["MigrationScript"],
        ) -> None:
            """Handle revision directives."""
            if autogenerate and cast("UpgradeOps", directives[0].upgrade_ops).is_empty():
                console.rule(
                    "[magenta]The generation of a migration file is being skipped because it would result in an empty file.",
                    style="magenta",
                    align="left",
                )
                console.rule(
                    "[magenta]More information can be found here. https://alembic.sqlalchemy.org/en/latest/autogenerate.html#what-does-autogenerate-detect-and-what-does-it-not-detect",
                    style="magenta",
                    align="left",
                )
                console.rule(
                    "[magenta]If you intend to create an empty migration file, use the --no-autogenerate option.",
                    style="magenta",
                    align="left",
                )
                directives.clear()

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Starting database upgrade process[/]", align="left")
        if message is None:
            message = "autogenerated" if no_prompt else Prompt.ask("Please enter a message describing this revision")

        sqlalchemy_config = get_config_by_bind_key(ctx, bind_key)
        alembic_commands = AlembicCommands(sqlalchemy_config=sqlalchemy_config)
        alembic_commands.revision(
            message=message,
            autogenerate=autogenerate,
            sql=sql,
            head=head,
            splice=splice,
            branch_label=branch_label,
            version_path=version_path,
            rev_id=rev_id,
            process_revision_directives=process_revision_directives,  # type: ignore[arg-type]
        )

    @database_group.command(name="drop-all", help="Drop all tables from the database.")
    @bind_key_option
    @no_prompt_option
    def drop_all(bind_key: Optional[str], no_prompt: bool) -> None:  # pyright: ignore[reportUnusedFunction]
        """Drop all tables from the database."""
        from anyio import run
        from rich.prompt import Confirm

        from advanced_alchemy.alembic.utils import drop_all
        from advanced_alchemy.base import metadata_registry

        ctx = cast("click.Context", click.get_current_context())
        console.rule("[yellow]Dropping all tables from the database[/]", align="left")
        input_confirmed = no_prompt or Confirm.ask(
            "[bold red]Are you sure you want to drop all tables from the database?"
        )

        async def _drop_all(
            configs: "Sequence[Union[SQLAlchemyAsyncConfig, SQLAlchemySyncConfig]]",
        ) -> None:
            for config in configs:
                engine = config.get_engine()
                await drop_all(engine, config.alembic_config.version_table_name, metadata_registry.get(config.bind_key))

        if input_confirmed:
            configs = [get_config_by_bind_key(ctx, bind_key)] if bind_key is not None else ctx.obj["configs"]
            run(_drop_all, configs)

    @database_group.command(name="dump-data", help="Dump specified tables from the database to JSON files.")
    @bind_key_option
    @click.option(
        "--table",
        "table_names",
        help="Name of the table to dump. Multiple tables can be specified. Use '*' to dump all tables.",
        type=str,
        required=True,
        multiple=True,
    )
    @click.option(
        "--dir",
        "dump_dir",
        help="Directory to save the JSON files. Defaults to WORKDIR/fixtures",
        type=click.Path(path_type=Path),
        default=Path.cwd() / "fixtures",
        required=False,
    )
    def dump_table_data(bind_key: Optional[str], table_names: tuple[str, ...], dump_dir: Path) -> None:  # pyright: ignore[reportUnusedFunction]
        """Dump table data to JSON files."""
        from anyio import run
        from rich.prompt import Confirm

        from advanced_alchemy.alembic.utils import dump_tables
        from advanced_alchemy.base import metadata_registry, orm_registry

        ctx = cast("click.Context", click.get_current_context())
        all_tables = "*" in table_names

        if all_tables and not Confirm.ask(
            "[yellow bold]You have specified '*'. Are you sure you want to dump all tables from the database?",
        ):
            return console.rule("[red bold]No data was dumped.", style="red", align="left")

        async def _dump_tables() -> None:
            configs = [get_config_by_bind_key(ctx, bind_key)] if bind_key is not None else ctx.obj["configs"]
            for config in configs:
                target_tables = set(metadata_registry.get(config.bind_key).tables)

                if not all_tables:
                    for table_name in set(table_names) - target_tables:
                        console.rule(
                            f"[red bold]Skipping table '{table_name}' because it is not available in the default registry",
                            style="red",
                            align="left",
                        )
                    target_tables.intersection_update(table_names)
                else:
                    console.rule("[yellow bold]Dumping all tables", style="yellow", align="left")

                models = [
                    mapper.class_ for mapper in orm_registry.mappers if mapper.class_.__table__.name in target_tables
                ]
                await dump_tables(dump_dir, config.get_session(), models)
                console.rule("[green bold]Data dump complete", align="left")

        return run(_dump_tables)

    return database_group