"""Signature-first manifest generation.

The annotated function signature is the source of truth: a Pydantic model is
synthesized from it with ``create_model``, and ``model_json_schema()`` becomes
the manifest's ``inputSchema``. The ``uiSchema`` is collected from
``Annotated[..., UI(...)]`` markers. No ``BaseModel`` subclass is required for
the top-level arguments (nested *structured* types are still models).
"""

from __future__ import annotations

import inspect
import pathlib
import re
import types as _pytypes
from ast import literal_eval
from typing import Annotated, Any, Union, get_args, get_origin

from pydantic import BaseModel, Field, create_model
from pydantic_core import PydanticUndefined

from ._schema import SPECIAL_DIR_ARGS, resolve_ref
from ._types import PMPath, UI


def _substitute_paths(ann):
    """Recursively rewrite pathlib.Path / PurePath leaves to PMPath, through
    Optional / list / dict / Annotated *and into nested Pydantic models*. Lets
    authors keep writing plain ``Path`` (even inside nested models) while still
    getting the permissive, URI-safe, format=path type — so URIs like ``slpm://``
    survive validation instead of being coerced into a corrupted pathlib.Path."""
    if isinstance(ann, type) and issubclass(ann, pathlib.PurePath):
        return PMPath
    if isinstance(ann, type) and issubclass(ann, BaseModel):
        return _substitute_model_paths(ann)
    if hasattr(ann, "__metadata__"):  # Annotated[...]
        base = _substitute_paths(ann.__origin__)
        return Annotated[tuple([base, *ann.__metadata__])]
    origin = get_origin(ann)
    if origin is None:
        return ann
    args = tuple(_substitute_paths(a) for a in get_args(ann))
    if origin in (Union, _pytypes.UnionType):
        return Union[args]
    return origin[args[0]] if len(args) == 1 else origin[args]


def _fresh_field(finfo):
    """A new FieldInfo carrying the original's key settings (default, description,
    discriminator, ...) but *not* its stored annotation/metadata — so it doesn't
    clobber the substituted annotation's own metadata (e.g. PMPath's format:path).
    Field-level constraints aren't carried (path fields rarely have them; PMPath
    already supplies its own), but defaults, docs and the discriminator are."""
    kwargs = {}
    if finfo.default is not PydanticUndefined:
        kwargs["default"] = finfo.default
    if finfo.default_factory is not None:
        kwargs["default_factory"] = finfo.default_factory
    for attr in ("description", "title", "alias", "discriminator"):
        value = getattr(finfo, attr, None)
        if value is not None:
            kwargs[attr] = value
    return Field(**kwargs)


# Rebuilt nested models are cached so a model used by several fields (or
# recursively) maps to a single rebuilt class.
_PATH_SUBBED_MODELS: dict = {}


def _substitute_model_paths(model):
    """Return ``model`` with every ``pathlib.Path`` field (at any depth) rewritten
    to :data:`PMPath`. Rebuilt via ``create_model(__base__=model, ...)`` so the
    original validators, config and untouched fields are inherited; only fields
    whose annotation actually changes are overridden. A model with no Path inside
    is returned unchanged."""
    if model in _PATH_SUBBED_MODELS:
        return _PATH_SUBBED_MODELS[model]
    # Reserve the slot first so a self-referential field terminates the recursion.
    _PATH_SUBBED_MODELS[model] = model

    overrides = {}
    for fname, finfo in model.model_fields.items():
        new_ann = _substitute_paths(finfo.annotation)
        if new_ann is not finfo.annotation:
            # Pair the substituted annotation with a *fresh* FieldInfo carrying
            # the original's key settings. Reusing the original FieldInfo clobbers
            # the new annotation's own metadata (PMPath's format:path is dropped
            # from the schema), so we copy only the relevant attributes across.
            overrides[fname] = (new_ann, _fresh_field(finfo))

    if not overrides:
        return model

    new_model = create_model(model.__name__, __base__=model, **overrides)
    _PATH_SUBBED_MODELS[model] = new_model
    return new_model


def _split_ui(annotation):
    """Return ``(annotation_for_pydantic, UI_marker_or_None)``, removing any
    ``UI`` marker from ``Annotated[...]`` so it never reaches Pydantic."""
    if hasattr(annotation, "__metadata__"):
        base = annotation.__origin__
        meta = list(annotation.__metadata__)
        ui = next((m for m in meta if isinstance(m, UI)), None)
        non_ui = [m for m in meta if not isinstance(m, UI)]
        pyd = Annotated[tuple([base, *non_ui])] if non_ui else base
        return pyd, ui
    return annotation, None


def _is_optional(ann) -> bool:
    """True if ``ann`` is a Union that already admits ``None`` (e.g. ``T | None``,
    ``Optional[T]``). Used to avoid double-wrapping a ``= None`` parameter that is
    already explicitly Optional."""
    return get_origin(ann) in (Union, _pytypes.UnionType) and type(None) in get_args(
        ann
    )


def _is_array_annotation(ann) -> bool:
    """True if ``ann`` is a list/tuple/set (optionally wrapped in Annotated or
    Optional). Used to place item-targeting widget hints under ``items``."""
    if hasattr(ann, "__metadata__"):
        ann = ann.__origin__
    origin = get_origin(ann)
    if origin in (Union, _pytypes.UnionType):
        return any(
            _is_array_annotation(a) for a in get_args(ann) if a is not type(None)
        )
    return origin in (list, tuple, set, frozenset)


def _ui_entry(ui: UI, is_array: bool = False) -> dict:
    entry: dict = {}
    if ui.hidden:
        # Hiding applies to the whole field, never the items.
        entry["ui:widget"] = "hidden"
    elif ui.widget:
        # A widget on an array targets each element -> nest under "items".
        if is_array:
            entry["items"] = {"ui:widget": ui.widget}
        else:
            entry["ui:widget"] = ui.widget
    opts: dict = {}
    if ui.group is not None:
        opts["group"] = ui.group
    if ui.enabled_by is not None:
        opts["enabledBy"] = ui.enabled_by
    if opts:
        entry["ui:options"] = opts
    if ui.title:
        entry["ui:title"] = ui.title
    if ui.help:
        entry["ui:help"] = ui.help
    if ui.labels:
        # For dict[...] (additionalProperties) maps: RJSF applies the uiSchema
        # under "additionalProperties" to each added entry. Label the key column
        # (custom "keyLabel", read by a WrapIfAdditionalTemplate) and use the
        # value label as the value field's title.
        ap: dict = {}
        if ui.labels.get("key"):
            ap["ui:options"] = {"keyLabel": ui.labels["key"]}
        if ui.labels.get("value"):
            ap["ui:title"] = ui.labels["value"]
        if ap:
            entry["additionalProperties"] = ap
    return entry


def _strip_titles(node):
    """Remove Pydantic's auto-generated ``title`` keys from schema nodes (only
    real schema nodes, never a property literally named 'title'). Labels come
    from ``UI(title=...)`` / ``description`` instead."""
    if isinstance(node, list):
        for sub in node:
            _strip_titles(sub)
        return node
    if not isinstance(node, dict):
        return node
    node.pop("title", None)
    for key in ("properties", "$defs"):
        if isinstance(node.get(key), dict):
            for sub in node[key].values():
                _strip_titles(sub)
    for key in ("items", "additionalProperties"):
        if isinstance(node.get(key), dict):
            _strip_titles(node[key])
    for key in ("anyOf", "oneOf", "allOf"):
        if isinstance(node.get(key), list):
            for sub in node[key]:
                _strip_titles(sub)
    return node


def _discriminator_hide_paths(schema):
    """Return uiSchema key-paths for every discriminated-union discriminator
    property (e.g. ``mode``). These should be hidden: the oneOf option selector
    already sets them, so showing them is redundant."""
    defs = schema.get("$defs", {})
    paths = []

    def walk(node, uipath, seen):
        node = resolve_ref(node, defs)
        if not isinstance(node, dict) or id(node) in seen:
            return
        seen = seen | {id(node)}
        disc = node.get("discriminator")
        if node.get("oneOf") and isinstance(disc, dict) and disc.get("propertyName"):
            paths.append(uipath + [disc["propertyName"]])
        for key in ("anyOf", "oneOf"):
            for branch in node.get(key, []) or []:
                walk(branch, uipath, seen)
        for name, sub in (node.get("properties") or {}).items():
            walk(sub, uipath + [name], seen)
        if isinstance(node.get("additionalProperties"), dict):
            walk(node["additionalProperties"], uipath + ["additionalProperties"], seen)
        if isinstance(node.get("items"), dict):
            walk(node["items"], uipath + ["items"], seen)

    walk(schema, [], frozenset())
    return paths


def _apply_hidden(ui, path):
    node = ui
    for key in path[:-1]:
        node = node.setdefault(key, {})
    node.setdefault(path[-1], {})["ui:widget"] = "hidden"


# First "." not followed by a file extension (so "file.smi" isn't a sentence end).
_FIRST_SENTENCE = re.compile(r"\.(?!smi|sdf)")


def _derive_label(name, description):
    """PlayMolecule/pmview label rule: the description's first sentence as the
    label when it is <= 4 words, otherwise a prettified field name."""
    if description:
        first = _FIRST_SENTENCE.split(description, 1)[0].strip()
        if first and len(first.split()) <= 4:
            return first[:1].upper() + first[1:]
    return (name[:1].upper() + name[1:]).replace("_", " ")


def _label_assignments(schema):
    """Yield (uiSchema-path, label) for every (incl. nested) field, per the rule
    above. The full description stays in the schema (rendered as help/tooltip)."""
    defs = schema.get("$defs", {})
    out = []

    def walk(node, uipath, seen):
        node = resolve_ref(node, defs)
        if not isinstance(node, dict) or id(node) in seen:
            return
        seen = seen | {id(node)}
        for key in ("anyOf", "oneOf"):
            for branch in node.get(key, []) or []:
                walk(branch, uipath, seen)
        for pname, sub in (node.get("properties") or {}).items():
            rsub = resolve_ref(sub, defs) if isinstance(sub, dict) else None
            desc = (sub.get("description") if isinstance(sub, dict) else None) or (
                rsub.get("description") if isinstance(rsub, dict) else None
            )
            out.append((uipath + [pname], _derive_label(pname, desc)))
            walk(sub, uipath + [pname], seen)
        if isinstance(node.get("additionalProperties"), dict):
            walk(node["additionalProperties"], uipath + ["additionalProperties"], seen)
        if isinstance(node.get("items"), dict):
            walk(node["items"], uipath + ["items"], seen)

    walk(schema, [], frozenset())
    return out


def _apply_path_widgets(schema, ui):
    """Give every ``format:"path"`` field the ``file`` widget (on ``items`` for
    arrays), so Path arguments render as file pickers without annotation. Output
    directories (outdir/scratchdir/execdir) are skipped."""
    defs = schema.get("$defs", {})

    def walk(node, uipath, seen):
        node = resolve_ref(node, defs)
        if not isinstance(node, dict) or id(node) in seen:
            return
        seen = seen | {id(node)}
        if node.get("format") == "path":
            if not (len(uipath) == 1 and uipath[0] in SPECIAL_DIR_ARGS):
                cur = ui
                for key in uipath:
                    cur = cur.setdefault(key, {})
                cur.setdefault("ui:widget", "file")  # never override an explicit widget
            return
        for key in ("anyOf", "oneOf"):
            for branch in node.get(key, []) or []:
                walk(branch, uipath, seen)
        for pname, sub in (node.get("properties") or {}).items():
            walk(sub, uipath + [pname], seen)
        if isinstance(node.get("additionalProperties"), dict):
            walk(node["additionalProperties"], uipath + ["additionalProperties"], seen)
        if isinstance(node.get("items"), dict):
            walk(node["items"], uipath + ["items"], seen)

    walk(schema, [], frozenset())


def _apply_label(ui, path, label):
    node = ui
    for key in path[:-1]:
        node = node.setdefault(key, {})
    leaf = node.setdefault(path[-1], {})
    if "ui:title" not in leaf:  # never override an explicit UI(title=...)
        leaf["ui:title"] = label


_DOC_NAME = re.compile(r"^(\S+)\s*:")
_DOC_CHOICES = re.compile(r"choices=([\(\[].*[\)\]])")
_DOC_GUI = re.compile(r"gui_options=(.*)")


def _parse_docstring(doc):
    """Parse a numpydoc ``Parameters`` section into
    ``{name: {doc, choices, gui_options}}``. Lets existing functions document
    params in the docstring (func2argparse style) instead of Annotated/Field."""
    out = {}
    if not doc:
        return out
    cur = None
    in_params = False
    for raw in doc.splitlines():
        line = raw.strip()
        if in_params:
            if _DOC_NAME.match(line):
                cur = _DOC_NAME.findall(line)[0]
                out[cur] = {"doc": ""}
                ch = _DOC_CHOICES.findall(line)
                if ch:
                    try:
                        out[cur]["choices"] = literal_eval(ch[0])
                    except Exception:
                        pass
                go = _DOC_GUI.findall(line)
                if go:
                    try:
                        out[cur]["gui_options"] = literal_eval(go[0])
                    except Exception:
                        pass
            elif cur is not None:
                out[cur]["doc"] += line + " "
        if line.startswith("Parameters"):
            in_params = True
        elif in_params and line == "":
            in_params = False
    return out


def _has_enum(node):
    if not isinstance(node, dict):
        return False
    if node.get("enum") is not None or node.get("const") is not None:
        return True
    return any(_has_enum(b) for b in node.get("anyOf", []) or [])


def _apply_docstring_compat(schema, ui, argdocs):
    """Fill in description / choices / gui_options from the docstring when the
    typed annotation didn't supply them, so unmodified functions still produce a
    good manifest. Annotation/Field always wins over the docstring."""
    props = schema.get("properties", {})
    for name, info in argdocs.items():
        if name not in props:
            continue
        node = props[name]
        doc = (info.get("doc") or "").strip()
        if doc and not node.get("description"):
            node["description"] = doc
        choices = info.get("choices")
        if choices and not _has_enum(node):
            node["enum"] = list(choices)
            node.setdefault("type", "string")
        go = info.get("gui_options") or {}
        if not go:
            continue
        ui_node = ui.setdefault(name, {})
        if go.get("group"):
            ui_node.setdefault("ui:options", {}).setdefault("group", go["group"])
        if go.get("enabled_by"):
            ui_node.setdefault("ui:options", {}).setdefault("enabledBy", go["enabled_by"])
        if go.get("label") and "ui:title" not in ui_node:
            ui_node["ui:title"] = go["label"]
        if go.get("hide"):
            ui_node["ui:widget"] = "hidden"
        cols = go.get("columns")
        if cols:  # legacy dict[str,str] table column names
            ap = ui_node.setdefault("additionalProperties", {})
            if len(cols) >= 1:
                ap.setdefault("ui:options", {}).setdefault("keyLabel", cols[0])
            if len(cols) >= 2 and "ui:title" not in ap:
                ap["ui:title"] = cols[1]


def signature_to_inputs_model(func):
    """Build ``(PydanticModel, uiSchema)`` from a function's annotated signature.

    Plain ``pathlib.Path`` annotations are auto-substituted with :data:`PMPath`.
    ``Annotated[T, UI(...)]`` markers are collected into the RJSF ``uiSchema``.
    Parameters whose name starts with ``_`` (and ``self``) are skipped.
    """
    sig = inspect.signature(func)
    # NB: inspect.get_annotations (not typing.get_type_hints) — on Python 3.10
    # get_type_hints implicitly wraps ``= None`` params in Optional[...], which
    # buries the Annotated UI marker inside a Union (and chokes hashing UI(labels=
    # {...})). inspect.get_annotations skips that rewrite on every version.
    hints = inspect.get_annotations(func, eval_str=True)

    fields: dict = {}
    ui: dict = {}
    order: list = []

    for name, param in sig.parameters.items():
        if name.startswith("_") or name == "self":
            continue
        ann = hints.get(name, Any)
        pyd_ann, ui_marker = _split_ui(ann)
        pyd_ann = _substitute_paths(pyd_ann)
        default = ... if param.default is inspect.Parameter.empty else param.default
        # A ``= None`` default means the parameter is Optional, so make the schema
        # nullable — otherwise the emitted ``default: null`` contradicts a bare
        # ``type`` and a sent ``null`` fails JSON-Schema validation. We do this
        # explicitly (rather than via typing.get_type_hints' implicit-Optional
        # rewrite) because the UI marker was already split out above, so wrapping
        # here doesn't bury it inside a Union.
        if default is None and not _is_optional(pyd_ann):
            pyd_ann = Union[pyd_ann, None]
        fields[name] = (pyd_ann, default)
        order.append(name)
        if ui_marker is not None:
            entry = _ui_entry(ui_marker, _is_array_annotation(pyd_ann))
            if entry:
                ui[name] = entry

    model = create_model(func.__name__ + "_inputs", **fields)
    ui["ui:order"] = order + ["*"]
    return model, ui


def func_to_entry(
    func,
    *,
    function: str | None = None,
    env: str = "base",
    resources: dict | None = None,
    outputs=None,
    tests: dict | None = None,
    description: str | None = None,
    strip_titles: bool = True,
):
    """Build a v2 manifest *function entry* from an annotated function.

    The non-argument fields (``function``, ``env``, ``resources``, ``outputs``,
    ``tests``) are carried verbatim, identical to v1. ``schemaVersion: 2`` is
    stamped so a detached function dict is self-describing.

    For a full app manifest (the sibling ``manifest.json`` template merged with
    one or more of these entries), use :func:`func_to_manifest` instead.
    """
    model, ui = signature_to_inputs_model(func)
    schema = model.model_json_schema()
    if strip_titles:
        _strip_titles(schema)

    # Docstring fallback: fill description / choices / gui_options for params
    # that didn't supply them via Annotated/Field (keeps old functions working).
    _apply_docstring_compat(schema, ui, _parse_docstring(inspect.getdoc(func)))

    # Path fields render as file pickers (file widget; on items for arrays).
    _apply_path_widgets(schema, ui)

    # Hide discriminator properties (the oneOf selector sets them).
    for path in _discriminator_hide_paths(schema):
        _apply_hidden(ui, path)

    # Derive a human label for every field (full description stays for help).
    for path, label in _label_assignments(schema):
        _apply_label(ui, path, label)

    if description is None:
        doc = inspect.getdoc(func)
        description = doc.strip().splitlines()[0] if doc else ""

    return {
        "function": function or f"{func.__module__}.{func.__qualname__}",
        "env": env,
        "resources": resources or {},
        "description": description,
        "schemaVersion": 2,
        "inputSchema": schema,
        "uiSchema": ui,
        "outputs": outputs if outputs is not None else {},
        "tests": tests or {},
    }


def func_to_manifest(functions, file=None, *, strip_titles: bool = True) -> dict:
    """Build a v2 *app manifest* by merging generated function entries into the
    sibling ``manifest.json`` (or ``manifest.yaml``) template next to ``file``.

    Drop-in for ``func2argparse.func_to_manifest`` at the app level. The template
    supplies per-function metadata (``function``, ``env``, ``examples``,
    ``outputs``, ``resources``, ``tests``) and app-level metadata
    (``container_config``, ``datasets``, ``citations``, ``meta_keywords``, ...).
    For each function the schema-derived fields (``schemaVersion`` /
    ``inputSchema`` / ``uiSchema`` / ``description`` / ``name``) are generated via
    :func:`func_to_entry` and merged into its matching ``functions[]`` entry;
    every other template field (e.g. ``examples``) is preserved, and any legacy
    v1 ``params`` is dropped.
    """
    import json
    import os

    if not isinstance(functions, (list, tuple)):
        functions = [functions]

    manifest: dict = {}
    if file is not None:
        base = os.path.dirname(os.path.abspath(file))
        mpath = os.path.join(base, "manifest.json")
        ypath = os.path.join(base, "manifest.yaml")
        if os.path.exists(mpath):
            with open(mpath) as f:
                manifest = json.load(f)
        elif os.path.exists(ypath):
            import yaml

            with open(ypath) as f:
                manifest = yaml.safe_load(f)

    if "functions" not in manifest:
        raise RuntimeError(
            "func_to_manifest requires a sibling manifest.json with a 'functions' "
            "section listing each function's metadata."
        )

    for func in functions:
        name = func.__name__
        try:
            idx = next(
                i
                for i, ff in enumerate(manifest["functions"])
                if str(ff.get("function", "")).endswith(f".{name}")
            )
        except StopIteration:
            raise RuntimeError(
                f"Function {name!r} not found in the manifest.json 'functions' "
                "section. Please add an entry for it."
            )
        entry = manifest["functions"][idx]
        generated = func_to_entry(
            func,
            function=entry.get("function"),
            env=entry.get("env", "base"),
            resources=entry.get("resources"),
            outputs=entry.get("outputs"),
            tests=entry.get("tests"),
            strip_titles=strip_titles,
        )
        # Merge the regenerated, schema-derived fields over the template entry;
        # template-only fields (e.g. examples) are preserved.
        entry.update(generated)
        entry["name"] = name
        entry.pop("params", None)

    return manifest


def build_app_manifest(name: str, functions: list, **app_meta) -> dict:
    """Wrap one or more function entries in a v2 manifest envelope."""
    return {"manifestVersion": 2, "name": name, "functions": list(functions), **app_meta}
