# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

import ast
from typing import Set, Tuple
import os
import re
import unicodedata

from datatailr.scheduler.constants import BATCH_JOB_ARGUMENTS
from datatailr.errors import TaskError
from datatailr.tag import get_tag
from datatailr.utils import run_shell_command


def get_available_env_args():
    """
    Get the available environment variables for batch job arguments.

    This function retrieves the environment variables that match the keys defined in DATATAILR_BATCH_JOB_ARGUMENTS.

    Returns:
        dict: A dictionary of available environment variables for batch jobs.
    """
    available_args = {}
    for key, value in os.environ.items():
        arg_key = key.replace("DATATAILR_BATCH_ARG_", "").lower()
        if arg_key in BATCH_JOB_ARGUMENTS:
            available_args[arg_key] = value
    return available_args


def normalize_name(name: str) -> str:
    """
    Normalize a name by converting it to lowercase, removing non unicode characters, and replacing spaces with underscores.

    Args:
        name (str): The name to normalize.

    Returns:
        str: The normalized name.
    """
    name = unicodedata.normalize("NFKC", name).lower()
    return re.sub(r"[^0-9a-z]+", "-", name).strip("-")


def get_base_url() -> str:
    """
    Get the job scheduler URL from environment variables.
    Returns:
        str: The job scheduler URL or an empty string if not set.
    """
    hostname = get_tag("public_hostname", cached=True)
    domain = get_tag("public_domain", cached=True)
    if hostname and domain:
        return f"https://{hostname}.{domain}"
    return ""


class RepoValidationError(TaskError):
    def __init__(self, message: str):
        super().__init__(message)
        self.message = message


def get_path_to_module(func_or_module) -> str:
    """
    Get the file system path to the *top-level* module for a function or module.

    Args:
        func_or_module: A Python function or module object whose top-level
                        package path should be resolved.

    Returns:
        str: The file system path to the top-level importable package.
    """
    import importlib.util
    import types

    # Determine the module name from the provided object.
    if isinstance(func_or_module, types.ModuleType):
        module_name = func_or_module.__name__
    else:
        module_name = getattr(func_or_module, "__module__", "")

    if not module_name:
        raise RepoValidationError(
            "Cannot determine module name for the provided object."
        )

    # We want the path to the top-level importable package for this object.
    # For example, if the function is imported as:
    #     from dashboards.dash.app import main
    # then func.__module__ is "dashboards.dash.app", but we want the filesystem
    # path corresponding to the top-level "dashboards" package, not the nested
    # "dash" submodule.
    top_level_module = module_name.split(".", 1)[0] if module_name else module_name

    try:
        # Resolve spec for the top-level package
        module_spec = importlib.util.find_spec(top_level_module)
    except ValueError:
        module_spec = None

    if module_spec:
        # Resolve the filesystem directory of the top package for module_name
        if module_spec.submodule_search_locations:
            # It's a package; get its directory
            package_root = os.path.abspath(
                next(iter(module_spec.submodule_search_locations))
            )
        elif module_spec.origin:
            # It's a module; use its containing directory
            package_root = os.path.abspath(os.path.dirname(module_spec.origin))
        else:
            raise RepoValidationError(f"Cannot resolve path for {module_name}")
    else:
        # Fallback to the underlying file path if we have it.
        if isinstance(func_or_module, types.ModuleType):
            file_path = getattr(func_or_module, "__file__", None)
        else:
            code_obj = getattr(func_or_module, "__code__", None)
            file_path = code_obj.co_filename if code_obj is not None else None

        if not file_path:
            raise RepoValidationError(
                f"Cannot resolve filesystem path for module {module_name}"
            )

        file_path = os.path.abspath(file_path)
        package_root = os.path.dirname(file_path)

    return package_root


def get_path_to_repo(path_to_module) -> str:
    """
    Get the path to the git repository root for the module where the function is defined.

    Args:
        func (function): The function whose repository path is to be retrieved.

    Returns:
        str: The file system path to the git repository root.
    """
    try:
        path_to_repo = run_shell_command(
            f"cd {path_to_module} && git rev-parse --show-toplevel"
        )[0]
    except RuntimeError:
        # simply return the parent directory if not a git repo
        path_to_repo = os.path.abspath(os.path.join(path_to_module, ".."))
    return path_to_repo


def get_imports_from_code(
    file_path: str, include_relatives: bool = False
) -> Tuple[Set[str], Set[str]]:
    """
    Extract import statements from a given code string.

    Args:
        file_path (str): The path to the Python file to analyze.
        include_relatives (bool): Whether to include relative imports in the results.
    Returns:
        Tuple[Set[str], Set[str]]: A tuple containing two sets:
            - The first set contains all imported module names (including relative imports if specified).
            - The second set contains the top-level package names of the imported modules.
    """
    if not os.path.isfile(file_path):
        raise FileNotFoundError(f"File not found: {file_path}")

    with open(file_path, "r", encoding="utf-8") as f:
        code = f.read()

    def _top_level_part(dotted: str) -> str:
        return dotted.split(".", 1)[0] if dotted else dotted

    try:
        tree = ast.parse(code)
    except SyntaxError as e:
        raise SystemExit(f"SyntaxError while parsing: {e}")

    modules: Set[str] = set()
    top_levels: Set[str] = set()

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                name = alias.name  # e.g., "pkg.sub"
                if name:
                    modules.add(name)
                    top_levels.add(_top_level_part(name))
        elif isinstance(node, ast.ImportFrom):
            # node.module may be None for `from . import x`
            base = node.module or ""
            if node.level and include_relatives:
                # Represent relative base with leading dots (e.g., "..utils")
                dotted = "." * node.level + base if base else "." * node.level
                modules.add(dotted)
                # No top-level package for relative-only imports
            if node.level == 0 and base:
                # Absolute import-from; we record the base (not the individual names)
                modules.add(base)
                top_levels.add(_top_level_part(base))

    return modules, top_levels
