import argparse
import ast
import inspect
import importlib.util
import importlib.metadata
import logging
import os
import subprocess
import sys
import sysconfig
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, List, Optional, Sequence, Set, Tuple

import requests


LOGGER = logging.getLogger("pyre_tools")
RESTART_MARKER = "PYRE_TOOLS_SHIELD_RESTARTED"

IGNORED_DIRS = {
    "venv",
    ".venv",
    "__pycache__",
    ".git",
    "env",
    ".idea",
    ".vscode",
    "node_modules",
    "dist",
    "build",
}
IGNORED_FILES = {"__init__.py", "setup.py"}
SUPPORTED_EXTENSIONS = {".py", ".pyw", ".ipynb"}

MODULE_ALIASES = {
    "cv2": "opencv-python",
    "pil": "pillow",
    "yaml": "pyyaml",
    "sklearn": "scikit-learn",
    "bs4": "beautifulsoup4",
    "crypto": "pycryptodome",
}


@dataclass(frozen=True)
class Requirement:
    name: str
    version: Optional[str]
    source: str
    origins: Tuple[str, ...]
    installed: bool

    @property
    def specifier(self) -> str:
        return f"{self.name}=={self.version}" if self.version else self.name


@dataclass(frozen=True)
class DependencyReport:
    target: str
    imports: Tuple[str, ...]
    requirements: Tuple[Requirement, ...]

    @property
    def missing_requirements(self) -> Tuple[Requirement, ...]:
        return tuple(requirement for requirement in self.requirements if not requirement.installed)


def configure_logging(verbose: bool = False) -> None:
    if not logging.getLogger().handlers:
        logging.basicConfig(level=logging.INFO, format="%(message)s")
    LOGGER.setLevel(logging.DEBUG if verbose else logging.INFO)


def extract_imports(file_path: os.PathLike[str] | str) -> Set[str]:
    path = Path(file_path)
    try:
        content = path.read_text(encoding="utf-8", errors="ignore")
        if not content.strip():
            return set()
        tree = ast.parse(content, filename=str(path))
    except Exception as exc:
        LOGGER.warning("Failed to parse %s: %s", path, exc)
        return set()

    imported_modules: Set[str] = set()
    ignored_dirs_lower = {name.lower() for name in IGNORED_DIRS}
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                imported_modules.add(alias.name.split(".")[0])
        elif isinstance(node, ast.ImportFrom):
            if node.module and node.level == 0:
                imported_modules.add(node.module.split(".")[0])
            elif node.level > 0:
                package_name = path.parent.name
                if package_name and package_name.lower() not in ignored_dirs_lower:
                    imported_modules.add(package_name)
    return imported_modules


def collect_python_files(path: os.PathLike[str] | str) -> List[Path]:
    base_path = Path(path)
    ignored_dirs_lower = {name.lower() for name in IGNORED_DIRS}
    ignored_files_lower = {name.lower() for name in IGNORED_FILES}
    collected: List[Path] = []
    for root, dirs, files in os.walk(base_path):
        dirs[:] = [directory for directory in dirs if directory.lower() not in ignored_dirs_lower]
        for file_name in files:
            extension = Path(file_name).suffix.lower()
            if extension in SUPPORTED_EXTENSIONS and file_name.lower() not in ignored_files_lower:
                collected.append(Path(root) / file_name)
    return collected


def collect_imports(path: os.PathLike[str] | str) -> Set[str]:
    base_path = Path(path)
    if base_path.is_file():
        return extract_imports(base_path)

    if base_path.is_dir():
        imports: Set[str] = set()
        for file_path in collect_python_files(base_path):
            imports.update(extract_imports(file_path))
        return imports

    raise FileNotFoundError(f"Path does not exist: {base_path}")


def get_installed_packages() -> tuple[dict[str, str], dict[str, str]]:
    installed: dict[str, str] = {}
    top_level_map: dict[str, str] = {}
    for dist in importlib.metadata.distributions():
        try:
            name = dist.metadata.get("Name")
            version = dist.version
            if not name or not version:
                continue

            key = name.lower()
            installed[key] = version

            top_level_text = dist.read_text("top_level.txt")
            if top_level_text:
                for line in top_level_text.splitlines():
                    top_level_name = line.strip()
                    if top_level_name:
                        top_level_map[top_level_name.lower()] = key
        except Exception:
            continue
    return installed, top_level_map


def is_stdlib(module_name: str) -> bool:
    if module_name in sys.builtin_module_names:
        return True
    try:
        spec = importlib.util.find_spec(module_name)
        if spec is None:
            return False
        if spec.origin in ("built-in", None):
            return True
        origin = str(spec.origin)
        if "frozen" in origin:
            return True
        normalized_origin = origin.lower().replace("\\", "/")
        if "/site-packages/" in normalized_origin or "/dist-packages/" in normalized_origin:
            return False
        stdlib_paths = [sysconfig.get_path("stdlib"), sysconfig.get_path("platstdlib")]
        return any(stdlib_path and origin.startswith(stdlib_path) for stdlib_path in stdlib_paths)
    except Exception:
        return False


def query_pypi_version(package_name: str, session: Optional[requests.Session] = None) -> Optional[str]:
    client = session or requests
    try:
        response = client.get(f"https://pypi.org/pypi/{package_name}/json", timeout=5)
        if response.status_code == 200:
            return response.json().get("info", {}).get("version")
    except Exception as exc:
        LOGGER.debug("PyPI lookup failed for %s: %s", package_name, exc)
    return None


def _resolve_requirement(
    module_name: str,
    installed_map: dict[str, str],
    top_level_map: dict[str, str],
    session: Optional[requests.Session] = None,
) -> Requirement:
    module_key = module_name.lower()
    resolved_name: Optional[str] = None
    resolved_version: Optional[str] = None
    source = "unknown"
    installed = False

    alias_name = MODULE_ALIASES.get(module_name) or MODULE_ALIASES.get(module_key)
    lookup_name = alias_name or module_name

    package_key = top_level_map.get(module_key)
    if package_key and package_key in installed_map:
        resolved_name = package_key
        resolved_version = installed_map[package_key]
        source = "installed(top_level)"
        installed = True
    elif module_key in installed_map:
        resolved_name = module_key
        resolved_version = installed_map[module_key]
        source = "installed(module_name)"
        installed = True
    else:
        pypi_version = query_pypi_version(lookup_name, session=session)
        if pypi_version:
            resolved_name = lookup_name
            resolved_version = pypi_version
            source = "pypi"

    if not resolved_name:
        resolved_name = lookup_name

    return Requirement(
        name=resolved_name,
        version=resolved_version,
        source=source,
        origins=(module_name,),
        installed=installed,
    )


def match_imports_to_requirements(imports: Iterable[str], include_stdlib: bool = False) -> List[Requirement]:
    installed_map, top_level_map = get_installed_packages()
    found: dict[str, Requirement] = {}
    session = requests.Session()

    for module_name in sorted(set(imports)):
        if not include_stdlib and is_stdlib(module_name):
            continue

        requirement = _resolve_requirement(module_name, installed_map, top_level_map, session=session)
        key = requirement.name.lower()
        existing = found.get(key)
        if existing:
            merged_origins = tuple(sorted(set(existing.origins) | {module_name}))
            preferred = existing if existing.installed or not requirement.installed else requirement
            found[key] = Requirement(
                name=preferred.name,
                version=preferred.version,
                source=preferred.source,
                origins=merged_origins,
                installed=preferred.installed,
            )
        else:
            found[key] = requirement

    return sorted(found.values(), key=lambda requirement: requirement.name.lower())


def analyze_path(path: os.PathLike[str] | str, include_stdlib: bool = False) -> DependencyReport:
    resolved_path = str(Path(path).resolve())
    imports = tuple(sorted(collect_imports(resolved_path)))
    requirements = tuple(match_imports_to_requirements(imports, include_stdlib=include_stdlib))
    return DependencyReport(target=resolved_path, imports=imports, requirements=requirements)


def find_missing_requirements(report_or_requirements: DependencyReport | Sequence[Requirement]) -> List[Requirement]:
    requirements = report_or_requirements.requirements if isinstance(report_or_requirements, DependencyReport) else report_or_requirements
    return [
        requirement
        for requirement in requirements
        if not requirement.installed and requirement.version is not None and requirement.source == "pypi"
    ]


def install_requirements(
    requirements: Sequence[Requirement],
    python_executable: Optional[str] = None,
    upgrade: bool = False,
) -> List[str]:
    installed_specs: List[str] = []
    executable = python_executable or sys.executable

    for requirement in find_missing_requirements(requirements):
        command = [executable, "-m", "pip", "install"]
        if upgrade:
            command.append("--upgrade")
        command.append(requirement.specifier)
        LOGGER.info("Installing %s", requirement.specifier)
        subprocess.check_call(command)
        installed_specs.append(requirement.specifier)

    return installed_specs


def _resolve_project_root(start_path: os.PathLike[str] | str) -> Path:
    path = Path(start_path).resolve()
    current = path.parent if path.is_file() else path
    markers = ("pyproject.toml", "setup.py", ".git")
    for candidate in (current, *current.parents):
        if any((candidate / marker).exists() for marker in markers):
            return candidate
    return current


def _is_pseudo_script_path(path: Path) -> bool:
    path_str = str(path)
    return path.name.startswith("<") and path.name.endswith(">") or path_str.startswith("<") and path_str.endswith(">")


def _resolve_caller_file() -> Path:
    current_file = Path(__file__).resolve()
    for frame_info in inspect.stack()[1:]:
        frame_path = Path(frame_info.filename).resolve()
        if frame_path != current_file and not _is_pseudo_script_path(frame_path):
            return frame_path
    argv_path = Path(sys.argv[0]).resolve()
    if not _is_pseudo_script_path(argv_path):
        return argv_path
    return Path.cwd()


def _restart_current_process() -> None:
    restart_env = dict(os.environ)
    restart_env[RESTART_MARKER] = "1"
    raise SystemExit(subprocess.call([sys.executable, *sys.argv], env=restart_env))


def _is_local_project_module(module_name: str, project_root: Path, caller_file: Path) -> bool:
    search_roots = [caller_file.parent, project_root]
    for root in search_roots:
        module_file = root / f"{module_name}.py"
        package_dir = root / module_name / "__init__.py"
        if module_file.exists() or package_dir.exists():
            return True
    return False


def _preflight_requirement(
    module_name: str,
    caller_file: Path,
    project_root: Path,
    installed_map: dict[str, str],
    top_level_map: dict[str, str],
) -> Requirement:
    module_key = module_name.lower()
    alias_name = MODULE_ALIASES.get(module_name) or MODULE_ALIASES.get(module_key)
    requirement_name = alias_name or module_name

    if _is_local_project_module(module_name, project_root, caller_file):
        return Requirement(
            name=module_name,
            version=None,
            source="local_module",
            origins=(module_name,),
            installed=True,
        )

    package_key = top_level_map.get(module_key)
    if package_key and package_key in installed_map:
        return Requirement(
            name=package_key,
            version=installed_map[package_key],
            source="installed(top_level)",
            origins=(module_name,),
            installed=True,
        )

    if module_key in installed_map:
        return Requirement(
            name=module_key,
            version=installed_map[module_key],
            source="installed(module_name)",
            origins=(module_name,),
            installed=True,
        )

    try:
        spec = importlib.util.find_spec(module_name)
        if spec is not None and not is_stdlib(module_name):
            return Requirement(
                name=requirement_name,
                version=None,
                source="importable",
                origins=(module_name,),
                installed=True,
            )
    except Exception:
        pass

    return Requirement(
        name=requirement_name,
        version=None,
        source="missing(preflight)",
        origins=(module_name,),
        installed=False,
    )


def preflight_current_script(
    caller_file: os.PathLike[str] | str,
    project_root: os.PathLike[str] | str,
    include_stdlib: bool = False,
) -> DependencyReport:
    caller_path = Path(caller_file).resolve()
    project_path = Path(project_root).resolve()
    if _is_pseudo_script_path(caller_path) or not caller_path.exists():
        return DependencyReport(target=str(project_path), imports=(), requirements=())
    imports = tuple(sorted(extract_imports(caller_path)))
    installed_map, top_level_map = get_installed_packages()
    requirements: List[Requirement] = []

    for module_name in imports:
        if not include_stdlib and is_stdlib(module_name):
            continue
        requirements.append(
            _preflight_requirement(
                module_name,
                caller_path,
                project_path,
                installed_map,
                top_level_map,
            )
        )

    return DependencyReport(
        target=str(caller_path),
        imports=imports,
        requirements=tuple(sorted(requirements, key=lambda requirement: requirement.name.lower())),
    )


def export_requirements(requirements: Sequence[Requirement], output_path: os.PathLike[str] | str) -> Path:
    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8") as handle:
        handle.write("# Auto-generated by pyre_tools\n")
        for requirement in requirements:
            origin_list = ", ".join(requirement.origins)
            handle.write(f"{requirement.specifier}    # {requirement.source} (from: {origin_list})\n")
    return output


def export_pyproject_toml(
    requirements: Sequence[Requirement],
    output_path: os.PathLike[str] | str,
    project_name: str,
) -> Path:
    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    python_requirement = f">={sys.version_info.major}.{sys.version_info.minor}"
    content = [
        "[project]",
        f'name = "{project_name}"',
        'version = "0.1.0"',
        'description = "Auto-generated by pyre_tools"',
        'readme = "README.md"',
        f'requires-python = "{python_requirement}"',
        "",
        "dependencies = [",
    ]
    for requirement in requirements:
        dependency = requirement.specifier if requirement.version else requirement.name
        content.append(f'    "{dependency}",')
    content.extend([
        "]",
        "",
        "[build-system]",
        'requires = ["setuptools>=61.0.0", "wheel"]',
        'build-backend = "setuptools.build_meta"',
    ])
    output.write_text("\n".join(content), encoding="utf-8")
    return output


def check_and_install_dependencies(
    path: os.PathLike[str] | str = ".",
    *,
    auto_install: bool = True,
    output_dir: Optional[os.PathLike[str] | str] = None,
    write_files: bool = False,
    include_stdlib: bool = False,
    upgrade: bool = False,
) -> DependencyReport:
    report = analyze_path(path, include_stdlib=include_stdlib)
    if write_files:
        target_path = Path(path)
        destination = Path(output_dir) if output_dir else (target_path.parent if target_path.is_file() else target_path)
        project_name = target_path.stem if target_path.is_file() else target_path.resolve().name or "project"
        export_requirements(report.requirements, destination / "requirements.txt")
        export_pyproject_toml(report.requirements, destination / "pyproject.toml", project_name)
    if auto_install:
        install_requirements(report.requirements, upgrade=upgrade)
    return report


def bootstrap(
    path: os.PathLike[str] | str = ".",
    *,
    auto_install: bool = True,
    output_dir: Optional[os.PathLike[str] | str] = None,
    write_files: bool = False,
    include_stdlib: bool = False,
    upgrade: bool = False,
) -> DependencyReport:
    return check_and_install_dependencies(
        path,
        auto_install=auto_install,
        output_dir=output_dir,
        write_files=write_files,
        include_stdlib=include_stdlib,
        upgrade=upgrade,
    )


def shield(
    path: Optional[os.PathLike[str] | str] = None,
    *,
    auto_install: bool = True,
    output_dir: Optional[os.PathLike[str] | str] = None,
    write_files: bool = False,
    include_stdlib: bool = False,
    upgrade: bool = False,
    restart_on_install: bool = True,
) -> DependencyReport:
    caller_file = _resolve_caller_file()
    if path is not None:
        target_path = Path(path).resolve()
    elif caller_file.is_file():
        target_path = _resolve_project_root(caller_file)
    else:
        target_path = Path.cwd()
    preflight_target = target_path if target_path.is_file() else caller_file
    project_root = target_path if target_path.is_dir() else _resolve_project_root(target_path)

    preflight_report = preflight_current_script(
        preflight_target,
        project_root,
        include_stdlib=include_stdlib,
    )
    if not preflight_report.missing_requirements:
        LOGGER.debug("Preflight passed for %s; skipping full project scan.", preflight_report.target)
        return preflight_report

    report = analyze_path(target_path, include_stdlib=include_stdlib)
    installed_specs: List[str] = []

    if write_files:
        destination = Path(output_dir) if output_dir else (target_path.parent if target_path.is_file() else target_path)
        project_name = target_path.stem if target_path.is_file() else target_path.resolve().name or "project"
        export_requirements(report.requirements, destination / "requirements.txt")
        export_pyproject_toml(report.requirements, destination / "pyproject.toml", project_name)

    if auto_install:
        installed_specs = install_requirements(report.requirements, upgrade=upgrade)

    if installed_specs and restart_on_install and os.environ.get(RESTART_MARKER) != "1":
        LOGGER.info("Restarting process after installing dependencies: %s", ", ".join(installed_specs))
        _restart_current_process()

    return report


def get_install_hint(package_name: str = "pyre_tools") -> str:
    return f"pip install {package_name}"


def format_module_not_found_hint(module_name: str = "pyre_tools", package_name: str = "pyre_tools") -> str:
    return f"ModuleNotFoundError: No module named '{module_name}'. Install it with: {get_install_hint(package_name)}"


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Analyze Python project imports and optionally install missing dependencies.",
        epilog=(
            "Module usage:\n"
            "  import pyre_tools; pyre_tools.shield()\n"
            "  from pyre_tools import bootstrap\n"
            "  report = bootstrap('path/to/project', auto_install=True)\n"
            "\n"
            "Hint:\n"
            "  Put `import pyre_tools; pyre_tools.shield()` at the top of your entry script."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("target", help="Python file or project directory to analyze")
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose logging")
    parser.add_argument("-o", "--output", help="Output directory for generated files")
    parser.add_argument("--write-files", action="store_true", help="Write requirements.txt and pyproject.toml")
    parser.add_argument(
        "--install",
        action="store_true",
        help="Install missing dependencies resolved from PyPI",
    )
    parser.add_argument(
        "--include-stdlib",
        action="store_true",
        help="Include standard library imports in the report",
    )
    parser.add_argument("--upgrade", action="store_true", help="Upgrade packages during installation")
    return parser


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)

    configure_logging(args.verbose)

    report = check_and_install_dependencies(
        args.target,
        auto_install=args.install,
        output_dir=args.output,
        write_files=args.write_files,
        include_stdlib=args.include_stdlib,
        upgrade=args.upgrade,
    )

    LOGGER.info("Target: %s", report.target)
    LOGGER.info("Imports: %s", ", ".join(report.imports) if report.imports else "(none)")
    if report.requirements:
        LOGGER.info("Requirements:")
        for requirement in report.requirements:
            suffix = " [installed]" if requirement.installed else ""
            LOGGER.info("  %s (%s)%s", requirement.specifier, requirement.source, suffix)
    else:
        LOGGER.info("Requirements: (none)")
    return 0
