from __future__ import annotations

import argparse
import json
from copy import deepcopy
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Type, Union

from samtranslator.compat import pydantic
from samtranslator.internal.schema_source import (
    any_cfn_resource,
    aws_serverless_api,
    aws_serverless_application,
    aws_serverless_connector,
    aws_serverless_function,
    aws_serverless_graphqlapi,
    aws_serverless_httpapi,
    aws_serverless_layerversion,
    aws_serverless_simpletable,
    aws_serverless_statemachine,
)
from samtranslator.internal.schema_source.common import BaseModel, LenientBaseModel


class Globals(BaseModel):
    Function: Optional[aws_serverless_function.Globals]
    Api: Optional[aws_serverless_api.Globals]
    HttpApi: Optional[aws_serverless_httpapi.Globals]
    SimpleTable: Optional[aws_serverless_simpletable.Globals]
    StateMachine: Optional[aws_serverless_statemachine.Globals]
    LayerVersion: Optional[aws_serverless_layerversion.Globals]


Resources = Union[
    aws_serverless_connector.Resource,
    aws_serverless_function.Resource,
    aws_serverless_simpletable.Resource,
    aws_serverless_statemachine.Resource,
    aws_serverless_layerversion.Resource,
    aws_serverless_api.Resource,
    aws_serverless_httpapi.Resource,
    aws_serverless_application.Resource,
    aws_serverless_graphqlapi.Resource,
]


class _ModelWithoutResources(LenientBaseModel):
    Globals: Optional[Globals]


class SamModel(_ModelWithoutResources):
    Resources: Dict[
        str,
        Union[
            Resources,
            # Ignore resources that are not AWS::Serverless::*
            any_cfn_resource.Resource,
        ],
    ]


class Model(_ModelWithoutResources):
    Resources: Dict[str, Resources]


def get_schema(model: Type[pydantic.BaseModel]) -> Dict[str, Any]:
    obj = model.schema()

    # http://json-schema.org/understanding-json-schema/reference/schema.html#schema
    # https://github.com/pydantic/pydantic/issues/1478
    # Validated in https://github.com/aws/serverless-application-model/blob/5c82f5d2ae95adabc9827398fba8ccfc3dbe101a/tests/schema/test_validate_schema.py#L91
    obj["$schema"] = "http://json-schema.org/draft-04/schema#"

    # Pydantic automatically adds title to model (https://github.com/pydantic/pydantic/issues/1051),
    # and the YAML extension for VS Code then shows 'PassThroughProp' as title for pass-through
    # properties (instead of the title of the property itself)... so manually deleting it.
    del obj["definitions"]["PassThroughProp"]["title"]

    return obj


def json_dumps(obj: Any) -> str:
    return json.dumps(obj, indent=2, sort_keys=True) + "\n"


def _replace_in_dict(d: Dict[str, Any], keyword: str, replace: Callable[[Dict[str, Any]], Any]) -> Dict[str, Any]:
    """
    Replace any dict containing keyword.

    replace() takes the containing dict as input, and returns its replacement.
    """
    if keyword in d:
        d = replace(d)
    for k, v in d.items():
        if isinstance(v, dict):
            d[k] = _replace_in_dict(v, keyword, replace)
    return d


def _deep_get(d: Dict[str, Any], path: List[str]) -> Dict[str, Any]:
    """
    Returns value at path defined by the keys in `path`.
    """
    for k in path:
        d = d[k]
    return d


def _add_embedded_connectors(schema: Dict[str, Any]) -> None:
    """
    Add embedded Connectors resource attribute to supported CloudFormation resources.
    """
    # We get the definition from an existing SAM resource
    embedded_connector = schema["definitions"][
        "samtranslator__internal__schema_source__aws_serverless_function__Resource"
    ]["properties"]["Connectors"]

    profiles = json.loads(Path("samtranslator/model/connector_profiles/profiles.json").read_text())

    # Only add the resource attributes to resources that support it
    source_resources = profiles["Permissions"].keys()
    for resource in source_resources:
        schema["definitions"][resource]["properties"]["Connectors"] = embedded_connector


def extend_with_cfn_schema(sam_schema: Dict[str, Any], cfn_schema: Dict[str, Any]) -> None:
    """
    Add CloudFormation resources and template syntax to SAM schema.
    """

    sam_defs = sam_schema["definitions"]
    cfn_defs = cfn_schema["definitions"]

    sam_props = sam_schema["properties"]
    cfn_props = cfn_schema["properties"]

    # Add Resources from CloudFormation schema to SAM schema
    cfn_resources = cfn_props["Resources"]["patternProperties"]["^[a-zA-Z0-9]+$"]["anyOf"]
    sam_props["Resources"]["additionalProperties"]["anyOf"].extend(cfn_resources)

    # Add any other top-level properties from CloudFormation schema to SAM schema
    for k in cfn_props:
        if k not in sam_props:
            sam_props[k] = cfn_props[k]

    # Add definitions from CloudFormation schema to SAM schema
    for k in cfn_defs:
        if k in sam_defs:
            raise Exception(f"Key {k} already in SAM schema definitions")
        sam_defs[k] = cfn_defs[k]

    _add_embedded_connectors(sam_schema)

    # Inject CloudFormation documentation to SAM pass-through properties
    def replace_passthrough(d: Dict[str, Any]) -> Dict[str, Any]:
        passthrough = d["__samPassThrough"]
        schema = deepcopy(_deep_get(cfn_schema, passthrough["schemaPath"]))
        schema["markdownDescription"] = passthrough["markdownDescriptionOverride"]
        schema["title"] = d["title"]  # Still want the original title, CFN property name could be different
        return schema

    _replace_in_dict(
        sam_schema,
        "__samPassThrough",
        replace_passthrough,
    )

    # The unified schema should include all supported properties
    sam_schema["additionalProperties"] = False


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--cfn-schema", help="input CloudFormation schema", type=Path, required=True)
    parser.add_argument("--sam-schema", help="output SAM schema", type=Path, required=True)
    parser.add_argument("--unified-schema", help="output unified schema", type=Path, required=True)
    args = parser.parse_args()

    sam_schema = get_schema(SamModel)
    args.sam_schema.write_text(json_dumps(sam_schema))

    unified_schema = get_schema(Model)
    cfn_schema = json.loads(args.cfn_schema.read_text())
    extend_with_cfn_schema(unified_schema, cfn_schema)
    args.unified_schema.write_text(json_dumps(unified_schema))


if __name__ == "__main__":
    main()