"""DAG file discovery and import."""
from __future__ import annotations

import importlib.util
import logging
import sys
from pathlib import Path

from crondag.sdk.dag import DAG, _dag_registry, clear_dag_registry

logger = logging.getLogger(__name__)


class DagBag:
    """Discovers and imports DAG files from a directory."""

    def __init__(self, dags_dir: str = "./dags"):
        self.dags_dir = Path(dags_dir).resolve()
        self.dags: dict[str, DAG] = {}
        self.import_errors: dict[str, str] = {}

    def collect_dags(self) -> dict[str, DAG]:
        """Scan dags_dir for Python files and import them."""
        self.dags.clear()
        self.import_errors.clear()

        if not self.dags_dir.exists():
            logger.warning("DAGs directory does not exist: %s", self.dags_dir)
            return self.dags

        # Clear the global registry before re-importing
        clear_dag_registry()

        for py_file in sorted(self.dags_dir.glob("*.py")):
            if py_file.name.startswith("_"):
                continue
            self._import_dag_file(py_file)

        # Collect all registered DAGs
        from crondag.sdk.dag import get_all_dags
        self.dags = get_all_dags()
        logger.info("Collected %d DAGs from %s", len(self.dags), self.dags_dir)
        return self.dags

    def _import_dag_file(self, path: Path):
        """Import a single DAG file."""
        module_name = f"crondag_user_dags.{path.stem}"
        try:
            spec = importlib.util.spec_from_file_location(module_name, path)
            if spec is None or spec.loader is None:
                self.import_errors[str(path)] = "Could not create module spec"
                return

            module = importlib.util.module_from_spec(spec)
            sys.modules[module_name] = module
            spec.loader.exec_module(module)
            logger.debug("Imported DAG file: %s", path)
        except Exception as e:
            self.import_errors[str(path)] = str(e)
            logger.error("Failed to import %s: %s", path, e)
        finally:
            # Clean up the module from sys.modules
            sys.modules.pop(module_name, None)
