from __future__ import annotations

import ast
from dataclasses import dataclass
import os
from pathlib import Path
from typing import List, Optional

DEFAULT_BOOK_FILENAMES = ("book.yaml", "book.yml", "iac.book.yaml", "iac.book.yml")
DEFAULT_SCRIPT_EXTENSIONS = ("", ".sh", ".bash")


def parse_book_dirs(env_value: Optional[str]) -> List[Path]:
    dirs: List[Path] = []
    if env_value:
        for part in env_value.split(":"):
            part = part.strip()
            if part:
                dirs.append(Path(part).expanduser())
    dirs.extend(
        [
            Path("./books"),
            Path("./.iac/books"),
            Path("~/.iac/books").expanduser(),
            Path("."),
        ]
    )

    out: List[Path] = []
    seen = set()
    for d in dirs:
        p = d.resolve()
        if str(p) in seen:
            continue
        seen.add(str(p))
        out.append(p)
    return out


def book_id_for_path(path: Path) -> str:
    """
    Derive a stable "book id" used in refs:
      - /x/foo.yaml -> foo
      - /x/foo.yml  -> foo
      - /x/foo/book.yaml -> foo
      - /x/foo/iac.book.yaml -> foo
    """
    path = path.resolve()
    name = path.name.lower()
    if name in DEFAULT_BOOK_FILENAMES:
        # directory-style book
        return path.parent.name
    # file-style book
    if path.suffix.lower() in (".yaml", ".yml"):
        return path.stem
    return path.stem


def list_books(book_dirs: List[Path]) -> List[Path]:
    """
    Find books:
      - *.yaml/*.yml files
      - */book.yaml, */iac.book.yaml (one level deep)
    """
    found: List[Path] = []
    for d in book_dirs:
        if not d.exists() or not d.is_dir():
            continue

        # file-style books
        for ext in ("*.yaml", "*.yml"):
            for p in d.glob(ext):
                if p.is_file():
                    found.append(p.resolve())

        # directory-style books (one level deep)
        for sub in d.iterdir():
            if not sub.is_dir():
                continue
            for fn in DEFAULT_BOOK_FILENAMES:
                candidate = sub / fn
                if candidate.exists() and candidate.is_file():
                    found.append(candidate.resolve())
                    break

    # unique stable sort by derived id then path
    uniq = {str(p): p for p in found}
    books = list(uniq.values())
    return sorted(books, key=lambda p: (book_id_for_path(p).lower(), str(p).lower()))


def resolve_book(book_ref: str, book_dirs: List[Path]) -> Path:
    """
    Resolve a book reference:
      - If book_ref is an existing file -> use it
      - If book_ref is an existing directory -> look for book.yaml etc inside
      - Else search book_dirs for:
          * <book_ref>.yaml / .yml
          * <book_ref>/book.yaml (and other DEFAULT_BOOK_FILENAMES)
    """
    p = Path(book_ref).expanduser()

    # Direct file path
    if p.exists() and p.is_file():
        return p.resolve()

    # Direct directory path
    if p.exists() and p.is_dir():
        for fn in DEFAULT_BOOK_FILENAMES:
            candidate = p / fn
            if candidate.exists() and candidate.is_file():
                return candidate.resolve()
        raise FileNotFoundError(
            f"Book directory found but no book.yaml inside: {book_ref}"
        )

    # Search by name in dirs
    candidates = [book_ref, f"{book_ref}.yaml", f"{book_ref}.yml"]
    for d in book_dirs:
        # file-style
        for c in candidates:
            q = d / c
            if q.exists() and q.is_file():
                return q.resolve()

        # directory-style: d/book_ref/book.yaml
        sub = d / book_ref
        if sub.exists() and sub.is_dir():
            for fn in DEFAULT_BOOK_FILENAMES:
                candidate = sub / fn
                if candidate.exists() and candidate.is_file():
                    return candidate.resolve()

    raise FileNotFoundError(f"Book not found: {book_ref}")


def split_book_step(ref: str) -> tuple[str, str]:
    if "@" not in ref:
        raise ValueError("step reference must be in the form 'book@step'")
    book, step = ref.split("@", 1)
    book = book.strip()
    step = step.strip()
    if not book or not step:
        raise ValueError("step reference must be in the form 'book@step'")
    return book, step


def parse_script_dirs(env_value: Optional[str]) -> List[Path]:
    dirs: List[Path] = []
    if env_value:
        for part in env_value.split(":"):
            part = part.strip()
            if part:
                dirs.append(Path(part).expanduser())
    dirs.extend(
        [
            Path("./scripts"),
            Path("./.iac/scripts"),
            Path("~/.iac/scripts").expanduser(),
            Path("."),
        ]
    )

    out: List[Path] = []
    seen = set()
    for d in dirs:
        p = d.resolve()
        if str(p) in seen:
            continue
        seen.add(str(p))
        out.append(p)
    return out


def parse_module_dirs(env_value: Optional[str]) -> List[Path]:
    dirs: List[Path] = []
    if env_value:
        for part in env_value.split(":"):
            part = part.strip()
            if part:
                dirs.append(Path(part).expanduser())
    dirs.extend(
        [
            Path("./modules"),
            Path("./.iac/modules"),
            Path("~/.iac/modules").expanduser(),
        ]
    )

    out: List[Path] = []
    seen = set()
    for d in dirs:
        p = d.resolve()
        if str(p) in seen:
            continue
        seen.add(str(p))
        out.append(p)
    return out


def resolve_script(
    script_ref: str, script_dirs: List[Path], *, book_path: Optional[Path] = None
) -> Path:
    p = Path(script_ref).expanduser()

    if p.exists() and p.is_file():
        return p.resolve()

    if not p.is_absolute() and book_path is not None:
        candidate = (book_path.parent / p).resolve()
        if candidate.exists() and candidate.is_file():
            return candidate

    for d in script_dirs:
        for ext in DEFAULT_SCRIPT_EXTENSIONS:
            q = d / (script_ref + ext)
            if q.exists() and q.is_file():
                return q.resolve()

    raise FileNotFoundError(f"Script not found: {script_ref}")


def split_module_function(ref: str) -> tuple[str, str]:
    if "@" not in ref:
        raise ValueError("module reference must be in the form 'module@function'")
    module_name, fn_name = ref.split("@", 1)
    module_name = module_name.strip()
    fn_name = fn_name.strip()
    if not module_name or not fn_name:
        raise ValueError("module reference must be in the form 'module@function'")
    return module_name, fn_name


@dataclass(frozen=True)
class ModuleFunctionInfo:
    name: str
    doc: str = ""


@dataclass(frozen=True)
class ModuleInfo:
    module: str
    path: str
    functions: List[ModuleFunctionInfo]
    error: str = ""


def list_modules(module_dirs: List[Path]) -> List[ModuleInfo]:
    skip_dir_names = {
        ".git",
        ".hg",
        ".svn",
        ".venv",
        "venv",
        "__pycache__",
        ".mypy_cache",
        ".pytest_cache",
        "site-packages",
        "node_modules",
        ".tox",
        ".nox",
        "dist",
        "build",
    }
    found: dict[str, Path] = {}
    for d in module_dirs:
        if not d.exists() or not d.is_dir():
            continue
        for root, dirs, files in os.walk(d):
            dirs[:] = [
                x for x in dirs if x not in skip_dir_names and not x.startswith(".")
            ]
            for fn in files:
                if not fn.endswith(".py"):
                    continue
                p = Path(root) / fn
                if not p.is_file():
                    continue
                try:
                    rel = p.resolve().relative_to(d.resolve())
                except Exception:
                    continue
                parts = list(rel.parts)
                if parts[-1] == "__init__.py":
                    mod_parts = parts[:-1]
                else:
                    mod_parts = parts
                    mod_parts[-1] = Path(mod_parts[-1]).stem
                if not mod_parts:
                    continue
                if any(part in skip_dir_names for part in mod_parts):
                    continue
                module_name = ".".join(mod_parts)
                found.setdefault(module_name, p.resolve())
    out: List[ModuleInfo] = []
    for module_name, path in sorted(found.items(), key=lambda kv: kv[0].lower()):
        out.append(_module_info_from_path(module_name, path))
    return out


def resolve_module_file(module_name: str, module_dirs: List[Path]) -> Path:
    clean = module_name.strip().strip(".")
    if not clean:
        raise FileNotFoundError("Module not found: empty module name")
    rel = Path(*clean.split("."))
    for d in module_dirs:
        f1 = (d / rel).with_suffix(".py")
        if f1.exists() and f1.is_file():
            return f1.resolve()
        f2 = d / rel / "__init__.py"
        if f2.exists() and f2.is_file():
            return f2.resolve()
    raise FileNotFoundError(f"Module not found: {module_name}")


def _module_info_from_path(module_name: str, path: Path) -> ModuleInfo:
    try:
        src = path.read_text(encoding="utf-8")
        root = ast.parse(src, filename=str(path))
        funcs: List[ModuleFunctionInfo] = []
        for node in root.body:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if node.name.startswith("_"):
                    continue
                # Keep public callables broad, and surface docstrings for UX.
                doc = ast.get_docstring(node) or ""
                funcs.append(ModuleFunctionInfo(name=node.name, doc=doc))
        funcs.sort(key=lambda x: x.name.lower())
        return ModuleInfo(module=module_name, path=str(path), functions=funcs)
    except Exception as e:
        return ModuleInfo(
            module=module_name, path=str(path), functions=[], error=str(e)
        )
