from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Sequence

from alembic import command
from alembic.config import Config
from sqlalchemy.engine import make_url

# Import shared constants and utils
from .constants import ALEMBIC_INI_TEMPLATE, ALEMBIC_SCRIPT_TEMPLATE
from .utils import (
    get_database_url_from_env,
    is_async_url,
    build_engine,
    ensure_database_exists,
)

# ---------- Alembic env.py renderers ----------

def _render_env_py(packages: Sequence[str]) -> str:
    packages_list = ", ".join(repr(p) for p in packages)
    tpl = """# Alembic env.py generated by svc-infra
from __future__ import annotations
import os
import logging
from importlib import import_module
from typing import Iterable, List

from alembic import context
from sqlalchemy import engine_from_config, pool
from sqlalchemy import create_engine
from sqlalchemy.engine import Connection

# Load logging configuration from alembic.ini
config = context.config
if config.config_file_name is not None:
    import logging.config
    logging.config.fileConfig(config.config_file_name)
logger = logging.getLogger(__name__)

# Discover metadata from packages
DISCOVER_PACKAGES: List[str] = [__PACKAGES_LIST__]
ENV_DISCOVER = os.getenv("ALEMBIC_DISCOVER_PACKAGES")
if ENV_DISCOVER:
    DISCOVER_PACKAGES = [s.strip() for s in ENV_DISCOVER.split(',') if s.strip()]

def _collect_metadata() -> List[object]:
    metadata = []
    for pkg in DISCOVER_PACKAGES:
        try:
            mod = import_module(pkg)
        except Exception as e:  # pragma: no cover
            logger.debug("Failed to import %s: %s", pkg, e)
            continue
        # Common conventions
        for attr in ("metadata", "MetaData", "Base", "base"):
            obj = getattr(mod, attr, None)
            if obj is None:
                # try nested models submodule
                try:
                    sub = import_module(f"{pkg}.models")
                    obj = getattr(sub, attr, None)
                except Exception:
                    obj = None
            if obj is None:
                continue
            md = getattr(obj, "metadata", None) or obj
            # If it's a declarative Base, extract .metadata
            if hasattr(md, "tables") and hasattr(md, "schema"):  # rough check for MetaData
                metadata.append(md)
    return metadata

target_metadata = _collect_metadata()

# Determine URL: prefer env var DATABASE_URL else alembic.ini sqlalchemy.url
env_db_url = os.getenv("DATABASE_URL")
if env_db_url:
    config.set_main_option("sqlalchemy.url", env_db_url)


def run_migrations_offline() -> None:
    url = config.get_main_option("sqlalchemy.url")
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
        compare_type=True,
        compare_server_default=True,
        include_schemas=True,
    )
    with context.begin_transaction():
        context.run_migrations()


def run_migrations_online() -> None:
    connectable = engine_from_config(
        config.get_section(config.config_ini_section, {}),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )
    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            compare_type=True,
            compare_server_default=True,
            include_schemas=True,
        )
        with context.begin_transaction():
            context.run_migrations()

if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()
"""
    return tpl.replace("__PACKAGES_LIST__", packages_list)


def _render_env_py_async(packages: Sequence[str]) -> str:
    packages_list = ", ".join(repr(p) for p in packages)
    tpl = """# Alembic async env.py generated by svc-infra
from __future__ import annotations
import os
import logging
from importlib import import_module
from typing import List

from alembic import context
from sqlalchemy.ext.asyncio import create_async_engine

config = context.config
if config.config_file_name is not None:
    import logging.config
    logging.config.fileConfig(config.config_file_name)
logger = logging.getLogger(__name__)

DISCOVER_PACKAGES: List[str] = [__PACKAGES_LIST__]
ENV_DISCOVER = os.getenv("ALEMBIC_DISCOVER_PACKAGES")
if ENV_DISCOVER:
    DISCOVER_PACKAGES = [s.strip() for s in ENV_DISCOVER.split(',') if s.strip()]

def _collect_metadata() -> List[object]:
    from importlib import import_module
    metadata = []
    for pkg in DISCOVER_PACKAGES:
        try:
            mod = import_module(pkg)
        except Exception as e:
            logger.debug("Failed to import %s: %s", pkg, e)
            continue
        for attr in ("metadata", "MetaData", "Base", "base"):
            obj = getattr(mod, attr, None)
            if obj is None:
                try:
                    sub = import_module(f"{pkg}.models")
                    obj = getattr(sub, attr, None)
                except Exception:
                    obj = None
            if obj is None:
                continue
            md = getattr(obj, "metadata", None) or obj
            if hasattr(md, "tables") and hasattr(md, "schema"):
                metadata.append(md)
    return metadata

target_metadata = _collect_metadata()

env_db_url = os.getenv("DATABASE_URL")
if env_db_url:
    config.set_main_option("sqlalchemy.url", env_db_url)

def do_run_migrations(connection):
    context.configure(
        connection=connection,
        target_metadata=target_metadata,
        compare_type=True,
        compare_server_default=True,
        include_schemas=True,
    )
    with context.begin_transaction():
        context.run_migrations()

async def run_migrations_online() -> None:
    connectable = create_async_engine(config.get_main_option("sqlalchemy.url"))
    async with connectable.connect() as connection:
        await connection.run_sync(do_run_migrations)
    await connectable.dispose()

if context.is_offline_mode():
    raise SystemExit("Run offline migrations with a sync env.py or set offline to False.")
else:
    import asyncio as _asyncio
    _asyncio.run(run_migrations_online())
"""
    return tpl.replace("__PACKAGES_LIST__", packages_list)


# ---------- Alembic init ----------

def init_alembic(
    project_root: Path | str,
    *,
    script_location: str = "migrations",
    async_db: bool = False,
    discover_packages: Optional[Sequence[str]] = None,
    overwrite: bool = False,
) -> Path:
    """Initialize Alembic in the target project directory.

    - Creates alembic.ini (or overwrites if requested).
    - Creates migrations/ with env.py and versions/.
    - env.py will read DATABASE_URL from environment at runtime, and
      discover model metadata from the provided packages or ALEMBIC_DISCOVER_PACKAGES.

    Returns the Path to the created migrations directory.
    """
    root = Path(project_root).resolve()
    root.mkdir(parents=True, exist_ok=True)

    migrations_dir = root / script_location
    versions_dir = migrations_dir / "versions"

    # Create alembic.ini
    alembic_ini = root / "alembic.ini"
    sqlalchemy_url = os.getenv("DATABASE_URL", "")
    dialect_name = (make_url(sqlalchemy_url).get_backend_name() if sqlalchemy_url else "")
    ini_contents = ALEMBIC_INI_TEMPLATE.format(
        script_location=script_location,
        sqlalchemy_url=sqlalchemy_url,
        dialect_name=dialect_name,
    )
    if alembic_ini.exists() and not overwrite:
        pass
    else:
        alembic_ini.write_text(ini_contents, encoding="utf-8")

    # Create migrations structure
    migrations_dir.mkdir(parents=True, exist_ok=True)
    versions_dir.mkdir(parents=True, exist_ok=True)

    # Ensure a local Alembic revision template exists for mako
    script_template = migrations_dir / "script.py.mako"
    if overwrite or not script_template.exists():
        script_template.write_text(ALEMBIC_SCRIPT_TEMPLATE, encoding="utf-8")

    # Render env.py
    pkgs = list(discover_packages or [])
    if not pkgs:
        # Common default; user can override via ALEMBIC_DISCOVER_PACKAGES
        # Keep empty to avoid import errors by default
        pkgs = []

    env_py = _render_env_py_async(pkgs) if async_db else _render_env_py(pkgs)
    env_path = migrations_dir / "env.py"
    if env_path.exists() and not overwrite:
        pass
    else:
        env_path.write_text(env_py, encoding="utf-8")

    return migrations_dir


# ---------- Alembic command helpers ----------


def _build_alembic_config(project_root: Path | str, script_location: str = "migrations", database_url: Optional[str] = None) -> Config:
    root = Path(project_root).resolve()
    cfg_path = root / "alembic.ini"
    cfg = Config(str(cfg_path)) if cfg_path.exists() else Config()
    cfg.set_main_option("script_location", str((root / script_location).resolve()))
    db_url = database_url or get_database_url_from_env(required=False) or cfg.get_main_option("sqlalchemy.url")
    if db_url:
        cfg.set_main_option("sqlalchemy.url", db_url)
    # make Alembic interpret prepend_sys_path using os.pathsep (":" on *nix, ";" on Windows)
    cfg.set_main_option("path_separator", "os")
    # keep your convenience sys.path entry
    cfg.set_main_option("prepend_sys_path", str(root))
    return cfg


def revision(
    project_root: Path | str,
    message: str,
    *,
    autogenerate: bool = False,
    head: str | None = "head",
    branch_label: str | None = None,
    version_path: str | None = None,
    sql: bool = False,
) -> None:
    """Create a new Alembic revision.

    Uses alembic.ini under project_root (created by init_alembic). Respects the
    DATABASE_URL environment variable if set, which overrides sqlalchemy.url.
    When autogenerate=True, Alembic compares discovered model metadata to the
    current database to produce a migration script.

    Args:
        project_root: Directory containing alembic.ini and the migrations/ folder.
        message: Human-readable message for the revision.
        autogenerate: If True, autogenerate operations based on model diffs.
        head: Base revision to branch from (default "head").
        branch_label: Optional branch label for the new revision.
        version_path: Optional versions subfolder to place the revision in.
        sql: If True, don't write files; emit SQL to stdout instead.
    """
    cfg = _build_alembic_config(project_root)
    command.revision(
        cfg,
        message=message,
        autogenerate=autogenerate,
        head=head,
        branch_label=branch_label,
        version_path=version_path,
        sql=sql,
    )


def upgrade(project_root: Path | str, revision_target: str = "head") -> None:
    """Apply migrations forward to the specified revision.

    Args:
        project_root: Directory containing alembic.ini and migrations/.
        revision_target: Target revision identifier (e.g. "head", "ae1027a7acf").
    """
    cfg = _build_alembic_config(project_root)
    command.upgrade(cfg, revision_target)


def downgrade(project_root: Path | str, revision_target: str = "-1") -> None:
    """Revert migrations down to the specified revision or relative step.

    Args:
        project_root: Directory containing alembic.ini and migrations/.
        revision_target: Target revision identifier or relative step (e.g. "-1").
    """
    cfg = _build_alembic_config(project_root)
    command.downgrade(cfg, revision_target)


def current(project_root: Path | str, verbose: bool = False) -> None:
    """Print the current database revision(s).

    Args:
        project_root: Directory containing alembic.ini and migrations/.
        verbose: If True, include detailed revision information.
    """
    cfg = _build_alembic_config(project_root)
    command.current(cfg, verbose=verbose)


def history(project_root: Path | str, verbose: bool = False) -> None:
    """Show the migration history for this project.

    Args:
        project_root: Directory containing alembic.ini and migrations/.
        verbose: If True, include down revisions, timestamps, and messages.
    """
    cfg = _build_alembic_config(project_root)
    command.history(cfg, verbose=verbose)


def stamp(project_root: Path | str, revision_target: str = "head") -> None:
    """Set the current database revision without running migrations.

    Useful for marking an existing database as up-to-date.

    Args:
        project_root: Directory containing alembic.ini and migrations/.
        revision_target: Target revision identifier (e.g. "head").
    """
    cfg = _build_alembic_config(project_root)
    command.stamp(cfg, revision_target)


def merge_heads(project_root: Path | str, message: Optional[str] = None) -> None:
    """Create a merge revision that joins multiple migration heads.

    Args:
        project_root: Directory containing alembic.ini and migrations/.
        message: Optional message to use for the merge revision.
    """
    cfg = _build_alembic_config(project_root)
    command.merge(cfg, "heads", message=message)


# ---------- High-level convenience API ----------


@dataclass(frozen=True)
class DBSetupResult:
    project_root: Path
    migrations_dir: Path
    alembic_ini: Path


def init_database_structure(
    *,
    project_root: Path | str,
    discover_packages: Optional[Sequence[str]] = None,
    async_db: bool = False,
    overwrite: bool = False,
    create_db_if_missing: bool = True,
) -> DBSetupResult:
    """High-level one-shot setup.

    - Ensures database exists (best-effort by dialect) using DATABASE_URL.
    - Initializes Alembic scaffolding (alembic.ini and migrations/ env.py).
    """
    db_url = get_database_url_from_env(required=True)
    if create_db_if_missing:
        ensure_database_exists(db_url)
    mig_dir = init_alembic(
        project_root,
        async_db=async_db or (db_url and is_async_url(db_url)),
        discover_packages=discover_packages,
        overwrite=overwrite,
    )
    return DBSetupResult(
        project_root=Path(project_root).resolve(),
        migrations_dir=mig_dir,
        alembic_ini=Path(project_root).resolve() / "alembic.ini",
    )


__all__ = [
    # env helpers
    "get_database_url_from_env",
    "is_async_url",
    # engines and db bootstrap
    "build_engine",
    "ensure_database_exists",
    # alembic init and commands
    "init_alembic",
    "revision",
    "upgrade",
    "downgrade",
    "current",
    "history",
    "stamp",
    "merge_heads",
    # high-level
    "init_database_structure",
    "DBSetupResult",
]
