#!/usr/bin/env python3
"""
This module offers functionality to semantically enrich JSON Schemas
by using `x-jsonld-context` annotations pointing to JSON-LD context documents,
and also to build ready-to-use JSON-LD contexts from annotated JSON Schemas.

An example of an annotated JSON schema:

```yaml
"$schema": https://json-schema.org/draft/2020-12/schema
"x-jsonld-context": observation.context.jsonld
title: Observation
type: object
required:
  - featureOfInterest
  - hasResult
  - resultTime
properties:
  featureOfInterest:
    type: string
  hasResult:
    type: object
  resultTime:
    type: string
    format: date-time
  observationCollection:
    type: string
```

... and its linked `x-jsonld-context`:

```json
{
  "@context": {
    "@version": 1.1,
    "sosa": "http://www.w3.org/ns/sosa/",
    "featureOfInterest": "sosa:featureOfInterest",
    "hasResult": "sosa:hasResult",
    "resultTime": "sosa:resultTime"
  }
}
```

A `SchemaAnnotator` instance would then generate the following annotated JSON schema:

```yaml
$schema: https://json-schema.org/draft/2020-12/schema
title: Observation
type: object
required:
- featureOfInterest
- hasResult
- observationTime
properties:
  featureOfInterest:
    'x-jsonld-id': http://www.w3.org/ns/sosa/featureOfInterest
    type: string
  hasResult:
    'x-jsonld-id': http://www.w3.org/ns/sosa/hasResult
    type: object
  observationCollection:
    type: string
  observationTime:
    'x-jsonld-id': http://www.w3.org/ns/sosa/resultTime
    format: date-time
    type: string
```

This schema can then be referenced from other entities that follow it (e.g., by using
[FG-JSON](https://github.com/opengeospatial/ogc-feat-geo-json) "definedby" links).

A client can then build a full JSON-LD `@context` (by using a `ContextBuilder` instance)
and use it when parsing plain-JSON entities:

```json
{
  "@context": {
    "featureOfInterest": "http://www.w3.org/ns/sosa/featureOfInterest",
    "hasResult": "http://www.w3.org/ns/sosa/hasResult",
    "observationTime": "http://www.w3.org/ns/sosa/resultTime"
  }
}
```

A JSON schema can be in YAML or JSON format (the annotated schema will use the same format
as the input one).

JSON schemas need to follow some rules to work with this tool:

* No nested `properties` are allowed. If they are needed, they should be put in a different
schema, and a `$ref` to it used inside the appropriate property definition.
* `allOf`/`someOf` root properties can be used to import other schemas (as long as they
contain `$ref`s to them).

This module can be run as a script, both for schema annotation and for context generation.

To annotate a schema (that already contains a `x-jsonld-context` to a JSON-LD context resource):

```shell
python -m ogc.na.annotate_schema --file path/to/schema.file.yaml
```

This will generate a new `annotated` directory replicating the layout of the input file
path (`/annotated/path/to/schema.file.yaml` in this example).

JSON-LD contexts can be built by adding a `-c` flag:

```shell
python -m ogc.na.annotate_schema -c --file annotated/path/to/schema.file.yaml
```

The resulting context will be printed to the standard output.

"""

from __future__ import annotations

import argparse
import dataclasses
import json
import logging
import re
import sys
from operator import attrgetter
from pathlib import Path
from typing import Any, AnyStr, Callable, Sequence, Iterable
from urllib.parse import urlparse, urljoin

import jsonschema
import requests_cache

from ogc.na.util import is_url, load_yaml, LRUCache, dump_yaml, merge_contexts, merge_dicts

logger = logging.getLogger(__name__)

ANNOTATION_PREFIX = 'x-jsonld-'
ANNOTATION_CONTEXT = 'x-jsonld-context'
ANNOTATION_ID = 'x-jsonld-id'
ANNOTATION_PREFIXES = 'x-jsonld-prefixes'
ANNOTATION_EXTRA_TERMS = 'x-jsonld-extra-terms'

CURIE_TERMS = '@id', '@type', '@index'

context_term_cache = LRUCache(maxsize=20)
requests_session = requests_cache.CachedSession('ogc.na.annotate_schema', backend='memory', expire_after=180)


@dataclasses.dataclass
class AnnotatedSchema:
    source: str | Path
    is_json: bool
    schema: dict


@dataclasses.dataclass
class ReferencedSchema:
    location: str | Path
    fragment: str | None = None
    subschema: dict | None = None
    full_contents: dict | None = None
    chain: list = dataclasses.field(default_factory=list)
    ref: str | Path = None
    is_json: bool = False


@dataclasses.dataclass
class ResolvedContext:
    context: dict[str, Any]
    prefixes: dict[str, str]


class SchemaResolver:

    def __init__(self, working_directory=Path()):
        self.working_directory = working_directory.resolve()
        self._schema_cache: dict[str | Path, tuple[Any, bool]] = {}

    @staticmethod
    def _get_branch(schema: dict, ref: str):
        path = re.sub(r'^#?/?', '', ref).split('/')
        pointer = schema
        for item in path:
            if path:
                pointer = pointer[item]
        return pointer

    def _load_contents(self, s: str | Path) -> tuple[dict, bool]:
        contents, is_json = self._schema_cache.get(s, (None, False))
        if contents is None:
            contents, is_json = load_json_yaml(read_contents(s)[0])
            self._schema_cache[s] = contents, is_json
        return contents, is_json

    def resolve_ref(self, ref: str | Path, from_schema: ReferencedSchema | None = None) -> tuple[Path | str, str]:
        location = ref
        fragment = None
        if isinstance(location, str):
            s = location.split('#', 1)
            fragment = s[1] if len(s) > 1 else None
            location = s[0]
            if not location:
                return location, fragment
            if not is_url(location):
                location = Path(location)

        if isinstance(location, Path):
            if location.is_absolute():
                location = location.resolve()
            elif not from_schema:
                location = self.working_directory.joinpath(location).resolve()
            elif not isinstance(from_schema.location, Path):
                location = urljoin(from_schema.location, str(location))
            else:
                location = from_schema.location.resolve().parent.joinpath(location).resolve()

        if location is None:
            raise ValueError(f'Unexpected ref type {type(ref).__name__}')

        return location, fragment

    def resolve_schema(self, ref: str | Path, from_schema: ReferencedSchema | None = None) -> ReferencedSchema:
        chain = from_schema.chain + [from_schema] if from_schema else []
        try:
            schema_source, fragment = self.resolve_ref(ref, from_schema)

            if not schema_source:
                if not from_schema:
                    raise ValueError('Local ref provided without an anchor: ' + ref)
                return ReferencedSchema(location=from_schema.location,
                                        fragment=ref[1:],
                                        subschema=SchemaResolver._get_branch(from_schema.full_contents, ref),
                                        full_contents=from_schema.full_contents,
                                        chain=chain,
                                        ref=ref,
                                        is_json=from_schema.is_json)

            contents, is_json = self._load_contents(schema_source)
            if fragment:
                return ReferencedSchema(location=schema_source, fragment=fragment,
                                        subschema=SchemaResolver._get_branch(contents, fragment),
                                        full_contents=contents,
                                        chain=chain,
                                        ref=ref,
                                        is_json=is_json)
            else:
                return ReferencedSchema(location=schema_source,
                                        subschema=contents,
                                        full_contents=contents,
                                        chain=chain,
                                        ref=ref,
                                        is_json=is_json)
        except Exception as e:
            f = f" from {from_schema.location}" if from_schema else ''
            raise IOError(f"Error resolving reference {ref}{f}") from e


def read_contents(location: Path | str | None) -> tuple[AnyStr | bytes, str]:
    """
    Reads contents from a file or URL

    @param location: filename or URL to load
    @return: a tuple with the loaded data (str or bytes) and the base URL, if any
    """
    if not location:
        raise ValueError('A location must be provided')

    if isinstance(location, Path) or not is_url(location):
        fn = Path(location)
        base_url = None
        logger.info('Reading file contents from %s', fn)
        with open(fn) as f:
            contents = f.read()
    else:
        base_url = location
        r = requests_session.get(location)
        r.raise_for_status()
        contents = r.content

    return contents, base_url


def load_json_yaml(contents: str | bytes) -> tuple[Any, bool]:
    """
    Loads either a JSON or a YAML file

    :param contents: contents to load
    :return: a tuple with the loaded document, and whether the detected format was JSON (True) or YAML (False)
    """
    try:
        obj = json.loads(contents)
        is_json = True
    except ValueError:
        obj = load_yaml(content=contents)
        is_json = False

    return obj, is_json


def resolve_ref(ref: str, fn_from: str | Path | None = None, url_from: str | None = None,
                base_url: str | None = None) -> tuple[Path | None, str | None]:
    """
    Resolves a `$ref`
    :param ref: the `$ref` to resolve
    :param fn_from: original name of the file containing the `$ref` (when it is a file)
    :param url_from: original URL of the document containing the `$ref` (when it is a URL)
    :param base_url: base URL of the document containing the `$ref` (if any)
    :return: a tuple of (Path, str) with only one None entry (the Path if the resolved
    reference is a file, or the str if it is a URL)
    """

    base_url = base_url or url_from
    if is_url(ref):
        return None, ref
    elif base_url:
        return None, urljoin(base_url, ref)
    else:
        fn_from = fn_from if isinstance(fn_from, Path) else Path(fn_from)
        ref = (fn_from.resolve().parent / ref).resolve()
        return ref, None


def resolve_context(ctx: Path | str | dict | list, expand_uris=True) -> ResolvedContext | None:
    if not ctx:
        return None

    prefixes = {}
    def expand_uri(curie, ctx_stack):
        if not expand_uris or not ctx_stack or ':' not in curie:
            return curie
        prefix, localpart = curie.split(':', 1)
        for c in reversed(ctx_stack):
            if prefix in c:
                term_val = c[prefix]
                prefix_uri = term_val if isinstance(term_val, str) else term_val.get('@id')
                prefixes[prefix] = prefix_uri
                return f"{prefix_uri}{localpart}"
        return curie

    def resolve_prop(term_val, ctx_stack):
        if isinstance(term_val, str):
            return expand_uri(term_val, ctx_stack)
        elif not isinstance(term_val, dict):
            return term_val
        for curie_term in CURIE_TERMS:
            curie_val = term_val.get(curie_term)
            if isinstance(curie_val, str) and ':' in curie_val:
                term_val[curie_term] = expand_uri(curie_val, ctx_stack)
        term_ctx = term_val.get('@context')
        if term_ctx:
            term_val['@context'] = resolve_inner(term_ctx, ctx_stack)
        return term_val

    def resolve_inner(inner_ctx, ctx_stack=None) -> ResolvedContext | None:
        resolved = None
        if isinstance(inner_ctx, Path) or (isinstance(inner_ctx, str) and not is_url(inner_ctx)):
            resolved = resolve_context(load_yaml(filename=ctx).get('@context'))
        elif isinstance(inner_ctx, str):
            r = requests_session.get(ctx)
            r.raise_for_status()
            fetched = r.json().get('@context')
            if fetched:
                resolved = resolve_context(fetched)
        elif isinstance(inner_ctx, Sequence):
            resolved_ctx = {}
            inner_prefixes = {}
            for ctx_entry in inner_ctx:
                if isinstance(ctx_entry, dict):
                    ctx_entry = ctx_entry.get('@context')
                resolved_entry = resolve_context(ctx_entry)
                inner_prefixes.update(resolved_entry.prefixes)
                resolved = ResolvedContext(merge_dicts(resolved_entry.context, resolved_ctx), inner_prefixes)
        else:
            resolved = ResolvedContext(inner_ctx, {})

        if not resolved or not resolved.context:
            return resolved

        if not ctx_stack:
            ctx_stack = [resolved.context]
        else:
            ctx_stack = ctx_stack + [resolved.context]

        for term in resolved.context.keys():
            term_val = resolved.context[term]
            resolved.context[term] = resolve_prop(term_val, ctx_stack)

        return resolved

    resolved_inner = resolve_inner(ctx)
    for p, puri in resolved_inner.prefixes.items():
        if p not in prefixes:
            prefixes[p] = puri

    return ResolvedContext(context=resolved_inner.context, prefixes=prefixes)


def validate_schema(schema: Any):
    jsonschema.validators.validator_for(schema).check_schema(schema)


class SchemaAnnotator:
    """
    Builds a set of annotated JSON schemas from a collection of input schemas
    that have `x-jsonld-context`s to JSON-LD context documents.

    The results will be stored in the `schemas` property (a dictionary of
    schema-path-or-url -> AnnotatedSchema mappings).
    """

    def __init__(self, ref_mapper: Callable[[str, Any], str] | None = None):
        """
        :param fn: file path to load (root schema)
        :param url: URL to load (root schema)
        :follow_refs: whether to follow `$ref`s (otherwise just annotate the provided root schema)
        """
        self._schema_resolver = SchemaResolver()
        self._ref_mapper = ref_mapper

    def process_schema(self, location: Path | str | None,
                       default_context: str | Path | dict | None = None) -> AnnotatedSchema:
        resolved_schema = self._schema_resolver.resolve_schema(location)
        schema = resolved_schema.subschema

        try:
            if '$schema' in schema and all(x not in schema for x in ('schema', 'openapi')):
                validate_schema(schema)
        except jsonschema.exceptions.SchemaError as e:
            raise ValueError(f"Could not parse schema from {location}") from e

        context_fn = schema.get(ANNOTATION_CONTEXT)
        schema.pop(ANNOTATION_CONTEXT, None)

        context = {}
        prefixes = {}

        if default_context and (context_fn != default_context
                                or not (isinstance(context_fn, Path)
                                        and isinstance(default_context, Path)
                                        and default_context.resolve() == context_fn.resolve())):
            # Only load the provided context if it's different from the schema-referenced one
            resolved_default_context = resolve_context(default_context)
            context, prefixes = attrgetter('context', 'prefixes')(resolved_default_context)

        if context_fn:
            context_fn, fragment = self._schema_resolver.resolve_ref(context_fn, resolved_schema)
            schema_context = resolve_context(context_fn)

            context = merge_contexts(context, schema_context.context)
            prefixes = prefixes | schema_context.prefixes

        def find_prop_context(prop, context_stack) -> dict | None:
            for ctx in reversed(context_stack):
                if prop in ctx:
                    prop_ctx = ctx[prop]
                    if isinstance(prop_ctx, str):
                        return {'@id': prop_ctx}
                    elif '@id' not in prop_ctx:
                        raise ValueError(f'Missing @id for property {prop} in context {json.dumps(ctx, indent=2)}')
                    else:
                        return {k: v for k, v in prop_ctx.items() if k.startswith('@')}

        def process_properties(obj: dict, context_stack: list[dict[str, Any]], level) -> Iterable[str]:
            properties: dict[str, dict] = obj.get('properties') if obj else None
            if not properties:
                return ()
            if not isinstance(properties, dict):
                raise ValueError('"properties" must be a dictionary')

            used_terms = set()
            for prop in list(properties.keys()):
                prop_value = properties[prop]
                prop_ctx = find_prop_context(prop, context_stack)
                if prop_ctx:
                    used_terms.add(prop)
                    prop_schema_ctx = {f"{ANNOTATION_PREFIX}{k[1:]}": v for k, v in prop_ctx.items() if k != '@context'}

                    if not prop_value or prop_value is True:
                        properties[prop] = prop_schema_ctx
                    else:
                        for k, v in prop_schema_ctx.items():
                            prop_value.setdefault(k, v)

                if prop_ctx and '@context' in prop_ctx:
                    prop_context_stack = context_stack + [prop_ctx['@context']]
                else:
                    prop_context_stack = context_stack
                used_terms.update(process_subschema(prop_value, prop_context_stack, level))

            return used_terms

        def process_subschema(subschema, context_stack, level=1) -> Iterable[str]:
            if not subschema or not isinstance(subschema, dict):
                return ()

            if self._ref_mapper and '$ref' in subschema:
                subschema['$ref'] = self._ref_mapper(subschema['$ref'], subschema)

            used_terms = set()

            # Annotate oneOf, allOf, anyOf
            for p in ('oneOf', 'allOf', 'anyOf'):
                collection = subschema.get(p)
                if collection and isinstance(collection, list):
                    for entry in collection:
                        used_terms.update(process_subschema(entry, context_stack, level + 1))

            # Annotate main schema
            schema_type = subschema.get('type')
            if not schema_type and 'properties' in subschema:
                schema_type = 'object'

            if schema_type == 'object':
                used_terms.update(process_properties(subschema, context_stack, level + 1))
            elif schema_type == 'array':
                for k in ('prefixItems', 'items', 'contains'):
                    used_terms.update(process_subschema(subschema.get(k), context_stack, level + 1))

            # Annotate $defs
            for defs_prop in ('$defs', 'definitions'):
                defs_value = subschema.get(defs_prop)
                if isinstance(defs_value, dict):
                    for defs_entry in defs_value.values():
                        used_terms.update(process_subschema(defs_entry, context_stack))

            # Get prefixes
            for p, bu in subschema.get(ANNOTATION_PREFIXES, {}).items():
                if p not in prefixes:
                    prefixes[p] = bu

            if len(context_stack) == level and context_stack[-1]:
                extra_terms = {k: (v if isinstance(v, str)
                                   else {f"{ANNOTATION_PREFIX}{vk[1:]}": vv for vk, vv in v.items()})
                               for k, v in context_stack[-1].items() if k[0] != '@'
                               and k not in prefixes
                               and k not in used_terms }
                if extra_terms:
                    subschema.setdefault(ANNOTATION_EXTRA_TERMS, {}).update(extra_terms)

            return used_terms

        process_subschema(schema, [context])

        if prefixes:
            schema[ANNOTATION_PREFIXES] = prefixes

        return AnnotatedSchema(
            source=location,
            is_json=resolved_schema.is_json,
            schema=schema
        )


class ContextBuilder:
    """
    Builds a JSON-LD context from a set of annotated JSON schemas.
    """

    def __init__(self, location: Path | str = None,
                 compact: bool = True, ref_mapper: Callable[[str], str] | None = None):
        """
        :param fn: file to load the annotated schema from
        :param url: URL to load the annotated schema from
        """
        self.context = {'@context': {}}
        self._parsed_schemas: dict[str | Path, dict] = {}
        self._ref_mapper = ref_mapper

        self._resolver = SchemaResolver()

        self.location = location

        context = self._build_context(self.location, compact)
        self.context = {'@context': context}

    def _build_context(self, schema_location: str | Path,
                       compact: bool = True) -> dict:

        parsed = self._parsed_schemas.get(schema_location)
        if parsed:
            return parsed

        root_schema = self._resolver.resolve_schema(schema_location)

        prefixes = {}

        own_context = {}

        if prefixes:
            own_context.update(prefixes)

        def read_properties(subschema: dict, from_schema: ReferencedSchema,
                            property_chain: list) -> dict | None:
            if not isinstance(subschema, dict):
                return None
            if subschema.get('type', 'object') != 'object':
                return None
            subschema_context = {}
            for prop, prop_val in subschema.get('properties', {}).items():
                if not isinstance(prop_val, dict):
                    continue
                prop_context = {}
                for term, term_val in prop_val.items():
                    if term.startswith(ANNOTATION_PREFIX) and term != ANNOTATION_CONTEXT:
                        prop_context['@' + term[len(ANNOTATION_PREFIX):]] = term_val
                inner_context = process_subschema(prop_val, from_schema, property_chain + ['properties', prop])
                if inner_context:
                    prop_context['@context'] = inner_context
                if isinstance(prop_context.get('@id'), str):
                    subschema_context[prop] = prop_context
                elif inner_context:
                    subschema_context = merge_contexts(subschema_context, inner_context, from_schema, property_chain)

            return subschema_context

        def process_subschema(subschema, from_schema, property_chain=None) -> dict | None:

            if property_chain is None:
                property_chain = []

            if not isinstance(subschema, dict):
                return None

            if '$ref' in subschema:
                ref = subschema['$ref']
                if self._ref_mapper:
                    ref = self._ref_mapper(ref)
                referenced_schema = self._resolver.resolve_schema(ref, from_schema)
                subschema = referenced_schema.subschema
                from_schema = referenced_schema

            if not subschema:
                return None

            sub_context = read_properties(subschema, from_schema, property_chain) or {}

            for i in ('allOf', 'anyOf', 'oneOf'):
                l = subschema.get(i)
                if isinstance(l, list):
                    for idx, sub_subschema in enumerate(l):
                        sub_context = merge_contexts(sub_context,
                                                     process_subschema(sub_subschema,
                                                                       from_schema,
                                                                       property_chain + [f"{i}[{idx}]"]),
                                                     from_schema, property_chain)

            for i in ('prefixItems', 'items', 'contains'):
                l = subschema.get(i)
                if isinstance(l, dict):
                    items_ctx = process_subschema(l, from_schema, property_chain + [i])
                    sub_context = merge_contexts(sub_context, items_ctx, from_schema, property_chain)

            if ANNOTATION_EXTRA_TERMS in subschema:
                for extra_term, extra_term_context in subschema[ANNOTATION_EXTRA_TERMS].items():
                    if extra_term not in sub_context:
                        if isinstance(extra_term_context, str):
                            extra_term_context = {'@id': extra_term_context}
                        sub_context[extra_term] = extra_term_context

            if sub_context:
                fixed_sub_context = {}
                for prop, prop_ctx in sub_context.items():
                    if prop.startswith('@'):
                        continue
                    if '@id' in prop_ctx:
                        fixed_sub_context[prop] = prop_ctx
                    elif prop_ctx.get('@context'):
                        merge_contexts(fixed_sub_context, prop_ctx['@context'], from_schema, property_chain)
                sub_context = fixed_sub_context

            sub_prefixes = subschema.get(ANNOTATION_PREFIXES)
            if isinstance(sub_prefixes, dict):
                prefixes.update(sub_prefixes)

            return sub_context

        own_context = merge_contexts(own_context, process_subschema(root_schema.subschema, root_schema),
                                     root_schema)

        if compact:

            rev_prefixes = {v: k for k, v in prefixes.items()}

            def compact_uri(uri: str) -> str:
                if uri.startswith('@'):
                    # JSON-LD keyword
                    return uri
                parts = urlparse(uri)
                if parts.fragment:
                    pref, suf = uri.rsplit('#', 1)
                    pref += '#'
                elif len(parts.path) > 1:
                    pref, suf = uri.rsplit('/', 1)
                    pref += '/'
                else:
                    return uri

                if pref in rev_prefixes:
                    return f"{rev_prefixes[pref]}:{suf}"
                else:
                    return uri

            def compact_branch(branch, existing_terms):

                for term in list(branch.keys()):
                    if term[0] == '@':
                        # skip special terms
                        continue
                    if term in existing_terms and existing_terms[term] == branch[term]:
                        # same term exists in ancestor -> delete
                        del branch[term]

                for term in list(branch.keys()):
                    term_value = branch[term]
                    if isinstance(term_value, dict):
                        if len(term_value) == 1 and '@id' in term_value:
                            branch[term] = compact_uri(term_value['@id'])
                        elif '@context' in term_value:
                            compact_branch(term_value['@context'], {**existing_terms, **branch})

            compact_branch(own_context, {})

        self._parsed_schemas[schema_location] = own_context
        return own_context


def dump_annotated_schemas(annotator: SchemaAnnotator, subdir: Path | str = 'annotated',
                           root_dir: Path | str | None = None,
                           output_fn_transform: Callable[[Path], Path] | None = None) -> list[Path]:
    """
    Creates a "mirror" directory (named `annotated` by default) with the resulting
    schemas annotated by a `SchemaAnnotator`.

    :param annotator: a `SchemaAnnotator` with the annotated schemas to read
    :param subdir: a name for the mirror directory
    :param root_dir: root directory for computing relative paths to schemas
    :param output_fn_transform: optional callable to transform the output path
    """
    wd = (Path(root_dir) if root_dir else Path()).resolve()
    subdir = subdir if isinstance(subdir, Path) else Path(subdir)
    result = []
    for path, schema in annotator.schemas.items():

        if isinstance(path, Path):
            output_fn = path.resolve().relative_to(wd)
        else:
            parsed = urlparse(str(path))
            output_fn = parsed.path

        output_fn = subdir / output_fn
        if output_fn_transform:
            output_fn = output_fn_transform(output_fn)
        output_fn.parent.mkdir(parents=True, exist_ok=True)

        if schema.is_json:
            with open(output_fn, 'w') as f:
                json.dump(schema.schema, f, indent=2)
        else:
            dump_yaml(schema.schema, output_fn)

        result.append(output_fn)

    return result


def _main():
    parser = argparse.ArgumentParser(
        prog='JSON Schema @id injector'
    )

    parser.add_argument(
        'schema',
        help='Entrypoint JSON Schema (filename or URL)',
    )

    parser.add_argument(
        '--context',
        required=False,
        help='Manually provided JSON-LD context (filename or URL)',
    )

    parser.add_argument(
        '-c',
        '--build-context',
        help='Build JSON-LD context fron annotated schemas',
        action='store_true'
    )

    parser.add_argument(
        '-F',
        '--no-follow-refs',
        help='Do not follow $ref\'s',
        action='store_true'
    )

    parser.add_argument(
        '-o',
        '--output',
        help='Output directory where to put the annotated schemas',
        default='annotated'
    )

    parser.add_argument(
        '-b',
        '--context-batch',
        help="Write JSON-LD context to a file with the same name and .jsonld extension (implies --build-context)",
        action='store_true',
    )

    parser.add_argument(
        '--stdout',
        action='store_true',
        help='Dump schemas to stdout'
    )

    args = parser.parse_args()

    if not args.schema:
        print('Error: no file and no URL provided', file=sys.stderr)
        parser.print_usage(file=sys.stderr)
        sys.exit(2)

    if args.build_context or args.context_batch:
        ctx_builder = ContextBuilder(args.schema)
        if args.context_batch:
            fn = Path(args.file).with_suffix('.jsonld')
            with open(fn, 'w') as f:
                json.dump(ctx_builder.context, f, indent=2)
        else:
            print(json.dumps(ctx_builder.context, indent=2))
    else:
        annotator = SchemaAnnotator()
        annotated = annotator.process_schema(args.schema, args.context)
        if args.stdout:
            print(dump_yaml(annotated.schema))
        else:
            dump_annotated_schemas(annotator, args.output)


if __name__ == '__main__':
    logging.basicConfig(
        stream=sys.stderr,
        level=logging.INFO,
        format='%(asctime)s,%(msecs)d %(levelname)-5s [%(filename)s:%(lineno)d] %(message)s',
    )

    _main()
