#!/usr/bin/env python3
"""
Generate the JSON schema for augur subsample configuration files.

This script creates the augur subsample configuration schema file with inline
comments for better maintainability.
"""

import json
import sys
import os
from pathlib import Path
from typing import Any, Dict

# Add augur package to Python path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from augur.filter.arguments import descriptions
from augur.proximity import proximity_argument_descriptions
from augur.subsample import FILTER_SAMPLE_CONFIG, PROXIMAL_SAMPLE_CONFIG


OUTPUT = Path(__file__).parent.parent / "augur" / "data" / "schema-subsample-config.json"


def validate(schema_options: Dict[str, Any], config: Dict[str, Any]):
    schema_keys = set(schema_options.keys())
    config_keys = set(config.keys())

    if extra_in_schema := schema_keys - config_keys:
        raise Exception(f"Keys in schema but not in config: {sorted(extra_in_schema)}")

    if missing_in_schema := config_keys - schema_keys:
        raise Exception(f"Keys in config but not in schema: {sorted(missing_in_schema)}")

def _required_properties(options: Dict[str, Any]) -> Dict[str, Any]:
    """Extract properties marked with ``"required": True`` and return a dict
    with a ``"required"`` key suitable for unpacking into a JSON Schema object
    definition.  The ``"required"`` key is removed from each property in-place.
    """
    required = []
    for key, prop in options.items():
        if prop.pop("required", False):
            required.append(key)
    return {"required": required} if required else {}

def create_schema():
    # Initialize the schema with metadata.
    schema = {
        "_description": "This file is generated by devel/regenerate-subsample-schema. Do not edit manually - edit the script instead.",
        "$schema": "http://json-schema.org/draft-07/schema#",
        "$id": "https://nextstrain.org/schemas/augur/subsample-config/v1",
        "title": "Configuration file to be supplied to `augur subsample --config`",
        "type": "object",
        "additionalProperties": False,
        "required": ["samples"],
    }

    # Define default options that can be used at top-level or sample-level.
    default_options = {
        "exclude": {
            "oneOf": [
                {"type": "string", "format": "filepath"},
                {
                    "type": "array",
                    "items": {"type": "string", "format": "filepath"}
                }
            ],
            "description": descriptions["exclude"]
        },
        "exclude_all": {
            "type": "boolean",
            "description": descriptions["exclude_all"]
        },
        "exclude_ambiguous_dates_by": {
            "type": "string",
            "enum": ["any", "day", "month", "year"],
            "description": descriptions["exclude_ambiguous_dates_by"]
        },
        "exclude_where": {
            "oneOf": [
                {"type": "string"},
                {
                    "type": "array",
                    "items": {"type": "string"}
                }
            ],
            "description": descriptions["exclude_where"]
        },
        "include": {
            "oneOf": [
                {"type": "string", "format": "filepath"},
                {
                    "type": "array",
                    "items": {"type": "string", "format": "filepath"}
                }
            ],
            "description": descriptions["include"]
        },
        "include_where": {
            "oneOf": [
                {"type": "string"},
                {
                    "type": "array",
                    "items": {"type": "string"}
                }
            ],
            "description": descriptions["include_where"]
        },
        "min_date": {
            "type": ["string", "integer"],
            "description": descriptions["min_date"]
        },
        "max_date": {
            "type": ["string", "integer"],
            "description": descriptions["max_date"]
        },
        "min_length": {
            "type": "integer",
            "description": descriptions["min_length"]
        },
        "max_length": {
            "type": "integer",
            "description": descriptions["max_length"]
        },
        "exclude_invalid": {
            "type": "boolean",
            "description": descriptions["exclude_invalid"]
        },
        "non_nucleotide": {
            "type": "boolean",
            "description": descriptions["non_nucleotide_config"]
        },
        "query": {
            "type": "string",
            "description": descriptions["query"]
        },
        "query_columns": {
            "oneOf": [
                {"type": "string"},
                {
                    "type": "array",
                    "items": {"type": "string"}
                }
            ],
            "description": descriptions["query_columns"]
        }
    }

    # Define sample options as properties in the schema.
    # Note: these map directly to augur filter options¹, but that is an
    # implementation detail that may change in the future.
    # ¹ augur/subsample.py:FILTER_SAMPLE_CONFIG
    filter_sample_options = {
        **default_options,
        "context_sample": {
            "type": "string",
            "description": "Use the outputs from another sample as the inputs for this sample. Value must be a sample name.",
        },
        "drop_sample": {
            "type": "boolean",
            "description": "Drop this sample from the final output",
        },
        "group_by": {
            "oneOf": [
                {"type": "string"},
                {
                    "type": "array",
                    "items": {"type": "string"}
                }
            ],
            "description": descriptions["group_by"]
        },
        "group_by_weights": {
            "type": "string",
            "format": "filepath",
            "description": descriptions["group_by_weights"]
        },
        "probabilistic_sampling": {
            "type": "boolean",
            "description": descriptions["probabilistic_sampling"]
        },
        "sequences_per_group": {
            "type": "integer",
            "description": descriptions["sequences_per_group"]
        },
        "max_sequences": {
            "type": "integer",
            "description": descriptions["subsample_max_sequences"]
        }
    }

    validate(filter_sample_options, FILTER_SAMPLE_CONFIG)

    # Define sample proximity options as properties in the schema.
    # These map directly to augur proximity options¹.
    # ¹ augur/subsample.py:PROXIMAL_SAMPLE_CONFIG
    proximal_sample_options = {
        "method": {
            "type": "string",
            "enum": ["hamming"],
            "description": proximity_argument_descriptions["method"]
        },
        "focal_sample": {
            # Note: must be the _name_ of a sample, but the schema can't enforce this.
            "required": True,
            "type": "string",
            "description": proximity_argument_descriptions["focal_sequences"]
        },
        "context_sample": {
            # Note: must be the _name_ of a sample, but the schema can't enforce this.
            "type": "string",
            "description": proximity_argument_descriptions["context_sequences"]
        },
        "drop_sample": {
            "type": "boolean",
            "description": "Drop this sample from final outputs",
        },
        "k": {
            "type": "integer",
            "description": proximity_argument_descriptions["k"]
        },
        "max_distance": {
            "type": "integer",
            "description": proximity_argument_descriptions["max_distance"]
        },
        "ignore_missing_data": {
            "type": "string",
            "description": proximity_argument_descriptions["ignore_missing_data"]
        },
    }

    validate(proximal_sample_options, PROXIMAL_SAMPLE_CONFIG)

    # Add definitions for default and sample options.
    schema["$defs"] = {
        "defaultProperties": {
            "type": "object",
            "additionalProperties": False,
            "properties": default_options
        },
        "filterSampleProperties": {
            "type": "object",
            "additionalProperties": False,
            "properties": filter_sample_options,
            **_required_properties(filter_sample_options),
        },
        "proximalSampleProperties": {
            "type": "object",
            "additionalProperties": False,
            "properties": proximal_sample_options,
            **_required_properties(proximal_sample_options),
        }
    }

    # Define the schema structure.
    schema["properties"] = {
        "defaults": {"$ref": "#/$defs/defaultProperties"},
        "samples": {
            "type": "object",
            "minProperties": 1,
            "patternProperties": {
                # Sample names can be any non-empty string.
                "^.+$": {
                    "oneOf": [
                        {"$ref": "#/$defs/filterSampleProperties"},
                        {"$ref": "#/$defs/proximalSampleProperties"}
                    ]
                }
            }
        }
    }
    # Note: this leaves room for additional properties such as 'output' and
    # 'proximity_target' as briefly drafted in
    # <https://github.com/nextstrain/WNV/pull/97>.

    return schema


def main():
    # Generate the schema as a dictionary.
    schema = create_schema()

    # Write to the schema file.
    with open(OUTPUT, 'w') as f:
        json.dump(schema, f, indent=4)
    print(f"Schema generated successfully: {OUTPUT}")


if __name__ == "__main__":
    main()
