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,
    prepare_process_env,
    repair_alembic_state_if_needed
)

# ---------- 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 typing import List

from alembic import context
from sqlalchemy import engine_from_config, pool

# 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__)

# --- sys.path bootstrap for src-layout projects ---
import sys, pathlib
prepend = config.get_main_option("prepend_sys_path") or ""
if prepend:
    if prepend not in sys.path:
        sys.path.insert(0, prepend)
    src_path = pathlib.Path(prepend) / "src"
    if src_path.exists():
        s = str(src_path)
        if s not in sys.path:
            sys.path.insert(0, s)

# Discover metadata from packages, but prefer the central ModelBase first
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]:
    # 0) Prefer svc_infra.db.base.ModelBase.metadata if available
    try:
        from svc_infra.db.base import ModelBase  # type: ignore
        md = getattr(ModelBase, "metadata", None)
        if md is not None and hasattr(md, "tables") and md.tables:
            return [md]
    except Exception as e:
        logger.debug("ModelBase not available or empty: %s", e)

    # 1) Fallback: discover across packages (your existing logic)
    import importlib, pkgutil, pathlib
    found: list[object] = []

    def _maybe_add(obj: object) -> None:
        md = getattr(obj, "metadata", None) or obj
        if hasattr(md, "tables") and hasattr(md, "schema"):
            found.append(md)

    pkgs = list(DISCOVER_PACKAGES)

    if not pkgs:
        roots = []
        if prepend:
            roots.append(pathlib.Path(prepend))
            roots.append(pathlib.Path(prepend) / "src")
        for root in roots:
            if not root or not root.exists():
                continue
            for p in root.iterdir():
                if p.is_dir() and (p / "__init__.py").exists():
                    pkgs.append(p.name)

    for pkg_name in pkgs:
        try:
            pkg = importlib.import_module(pkg_name)
        except Exception as e:
            logger.debug("Failed to import %s: %s", pkg_name, e)
            continue

        for attr in ("metadata", "MetaData", "Base", "base"):
            obj = getattr(pkg, attr, None)
            if obj is not None:
                _maybe_add(obj)

        for subname in ("models",):
            try:
                sub = importlib.import_module(f"{pkg_name}.{subname}")
                for attr in ("metadata", "MetaData", "Base", "base"):
                    obj = getattr(sub, attr, None)
                    if obj is not None:
                        _maybe_add(obj)
            except Exception:
                pass

        mod_path = getattr(pkg, "__path__", None)
        if not mod_path:
            continue
        for _, name, ispkg in pkgutil.walk_packages(mod_path, prefix=pkg_name + "."):
            if ispkg or not any(x in name for x in (".models", ".db", ".orm", ".entities")):
                continue
            try:
                mod = importlib.import_module(name)
            except Exception:
                continue
            for attr in ("metadata", "MetaData", "Base", "base"):
                obj = getattr(mod, attr, None)
                if obj is not None:
                    _maybe_add(obj)

    uniq, seen = [], set()
    for md in found:
        if id(md) not in seen:
            seen.add(id(md))
            uniq.append(md)
    return uniq

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 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 typing import List

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

# 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__)

# --- sys.path bootstrap for src-layout projects ---
import sys, pathlib
prepend = config.get_main_option("prepend_sys_path") or ""
if prepend:
    if prepend not in sys.path:
        sys.path.insert(0, prepend)
    src_path = pathlib.Path(prepend) / "src"
    if src_path.exists():
        s = str(src_path)
        if s not in sys.path:
            sys.path.insert(0, s)

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]:
    # 0) Prefer ModelBase.metadata
    try:
        from svc_infra.db.base import ModelBase  # type: ignore
        md = getattr(ModelBase, "metadata", None)
        if md is not None and hasattr(md, "tables") and md.tables:
            return [md]
    except Exception as e:
        logger.debug("ModelBase not available or empty: %s", e)

    # 1) Fallback discovery (same as sync)
    import importlib, pkgutil, pathlib
    found: list[object] = []

    def _maybe_add(obj: object) -> None:
        md = getattr(obj, "metadata", None) or obj
        if hasattr(md, "tables") and hasattr(md, "schema"):
            found.append(md)

    pkgs = list(DISCOVER_PACKAGES)
    if not pkgs:
        roots = []
        if prepend:
            roots.append(pathlib.Path(prepend))
            roots.append(pathlib.Path(prepend) / "src")
        for root in roots:
            if not root or not root.exists():
                continue
            for p in root.iterdir():
                if p.is_dir() and (p / "__init__.py").exists():
                    pkgs.append(p.name)

    for pkg_name in pkgs:
        try:
            pkg = importlib.import_module(pkg_name)
        except Exception as e:
            logger.debug("Failed to import %s: %s", pkg_name, e)
            continue

        for attr in ("metadata", "MetaData", "Base", "base"):
            obj = getattr(pkg, attr, None)
            if obj is not None:
                _maybe_add(obj)

        for subname in ("models",):
            try:
                sub = importlib.import_module(f"{pkg_name}.{subname}")
                for attr in ("metadata", "MetaData", "Base", "base"):
                    obj = getattr(sub, attr, None)
                    if obj is not None:
                        _maybe_add(obj)
            except Exception:
                pass

        mod_path = getattr(pkg, "__path__", None)
        if not mod_path:
            continue
        for _, name, ispkg in pkgutil.walk_packages(mod_path, prefix=pkg_name + "."):
            if ispkg or not any(x in name for x in (".models", ".db", ".orm", ".entities")):
                continue
            try:
                mod = importlib.import_module(name)
            except Exception:
                continue
            for attr in ("metadata", "MetaData", "Base", "base"):
                obj = getattr(mod, attr, None)
                if obj is not None:
                    _maybe_add(obj)

    uniq, seen = [], set()
    for md in found:
        if id(md) not in seen:
            seen.add(id(md))
            uniq.append(md)
    return uniq

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:
    url = config.get_main_option("sqlalchemy.url")
    connectable = create_async_engine(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.ini + migrations/ scaffold.

    Example:
        >>> # DATABASE_URL from env; discovery via ModelBase or fallback scan
        >>> init_alembic(".", async_db=False, overwrite=False)

    Returns:
        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"

    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")

    migrations_dir.mkdir(parents=True, exist_ok=True)
    versions_dir.mkdir(parents=True, exist_ok=True)

    script_template = migrations_dir / "script.py.mako"
    need_template_write = overwrite or not script_template.exists()
    if not need_template_write and script_template.exists():
        try:
            current = script_template.read_text(encoding="utf-8")
            # If the template doesn't have the standard Mako slots, rewrite it.
            if ("${upgrades" not in current) or ("${downgrades" not in current):
                need_template_write = True
        except Exception:
            need_template_write = True

    if need_template_write:
        script_template.write_text(ALEMBIC_SCRIPT_TEMPLATE, encoding="utf-8")

    pkgs = list(discover_packages or [])
    if not pkgs:
        pkgs = []

    env_py_text = _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:
        try:
            existing = env_path.read_text(encoding="utf-8")
            if "DISCOVER_PACKAGES:" not in existing:
                env_path.write_text(env_py_text, encoding="utf-8")
        except Exception:
            env_path.write_text(env_py_text, encoding="utf-8")
    else:
        env_path.write_text(env_py_text, encoding="utf-8")

    return migrations_dir


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

def _build_alembic_config(
        project_root: Path | str,
        script_location: str = "migrations",
) -> Config:
    """Build Alembic Config based on files + environment."""
    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()))

    # Respect env if set (env.py also prefers env at runtime)
    env_db_url = os.getenv("DATABASE_URL", "").strip()
    if env_db_url:
        cfg.set_main_option("sqlalchemy.url", env_db_url)

    cfg.set_main_option("path_separator", "os")
    cfg.set_main_option("prepend_sys_path", str(root))
    return cfg

def _ensure_db_at_head(cfg: Config) -> None:
    """
    Bring the target database up to the current head if it's behind.
    Alembic raises 'Target database is not up to date' during autogenerate
    when the DB isn't at head; this preflight avoids that.
    """
    # This is idempotent if already at head.
    command.upgrade(cfg, "head")

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,
        ensure_head_before_autogenerate: bool = True,
) -> None:
    """
    Create a new Alembic revision.

    Example (autogenerate):
        >>> revision(".", "add orders", autogenerate=True)

    Requirements:
        - DATABASE_URL must be set in the environment.
        - Model discovery is automatic (prefers ModelBase.metadata).
    """
    prepare_process_env(project_root)  # no URL/pkgs
    cfg = _build_alembic_config(project_root)
    repair_alembic_state_if_needed(cfg)

    if autogenerate and ensure_head_before_autogenerate:
        if not (cfg.get_main_option("sqlalchemy.url") or os.getenv("DATABASE_URL")):
            raise RuntimeError("DATABASE_URL is not set.")
        _ensure_db_at_head(cfg)

    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.

    Example:
        >>> upgrade(".")          # to head
        >>> upgrade(".", "base")  # or to a specific rev
    """
    prepare_process_env(project_root)
    cfg = _build_alembic_config(project_root)
    repair_alembic_state_if_needed(cfg)
    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)
    repair_alembic_state_if_needed(cfg)
    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)
    repair_alembic_state_if_needed(cfg)
    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)
    repair_alembic_state_if_needed(cfg)
    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)
    repair_alembic_state_if_needed(cfg)
    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 SetupAndMigrateResult:
    """Structured outcome of setup_and_migrate."""
    project_root: Path
    migrations_dir: Path
    alembic_ini: Path
    created_initial_revision: bool
    created_followup_revision: bool
    upgraded: bool

def setup_and_migrate(
        *,
        project_root: Path | str,
        async_db: bool | None = None,
        overwrite_scaffold: bool = False,
        create_db_if_missing: bool = True,
        create_followup_revision: bool = True,
        initial_message: str = "initial schema",
        followup_message: str = "autogen",
) -> SetupAndMigrateResult:
    """
    Ensure DB + Alembic are ready and up-to-date.

    Examples:
        >>> # First run (DATABASE_URL already set in env)
        >>> setup_and_migrate(project_root=".")
        >>>
        >>> # Later, after editing models
        >>> setup_and_migrate(project_root=".", create_followup_revision=True)

    Notes:
        - Reads DATABASE_URL from environment. Does not set it.
        - Model discovery is automatic via env.py (prefers ModelBase.metadata).
    """
    root = Path(project_root).resolve()
    prepare_process_env(root)

    db_url = get_database_url_from_env(required=True)
    if create_db_if_missing:
        ensure_database_exists(db_url)

    from sqlalchemy.engine import make_url as _make_url
    is_async = async_db if async_db is not None else is_async_url(_make_url(db_url))

    mig_dir = init_alembic(
        root,
        async_db=is_async,
        discover_packages=None,   # rely on auto-discovery only
        overwrite=overwrite_scaffold,
    )
    versions_dir = mig_dir / "versions"
    alembic_ini = root / "alembic.ini"

    cfg = _build_alembic_config(project_root=root)
    repair_alembic_state_if_needed(cfg)

    created_initial = False
    created_followup = False
    upgraded = False

    try:
        upgrade(root)   # safe if nothing to do
        upgraded = True
    except Exception:
        pass

    def _has_revisions() -> bool:
        return any(versions_dir.glob("*.py"))

    if not _has_revisions():
        revision(
            project_root=root,
            message=initial_message,
            autogenerate=True,
            ensure_head_before_autogenerate=True,
        )
        created_initial = True
        upgrade(root)
        upgraded = True
    elif create_followup_revision:
        revision(
            project_root=root,
            message=followup_message,
            autogenerate=True,
            ensure_head_before_autogenerate=True,
        )
        created_followup = True
        upgrade(root)
        upgraded = True

    return SetupAndMigrateResult(
        project_root=root,
        migrations_dir=mig_dir,
        alembic_ini=alembic_ini,
        created_initial_revision=created_initial,
        created_followup_revision=created_followup,
        upgraded=upgraded,
    )


__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
    "setup_and_migrate",
    "SetupAndMigrateResult",
]
