"""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 create_model

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. Lets authors keep writing plain
    ``Path`` while still getting the permissive, URI-safe, format=path type."""
    if isinstance(ann, type) and issubclass(ann, pathlib.PurePath):
        return PMPath
    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 _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_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 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("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
        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_manifest(
    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.
    """
    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 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}
