"""ACDL Contract Resolver — resolve a consumer contract to a Target Stack instance.

The contract resolver is the bridge between the consumer's declared intent
(a contract YAML) and the platform's executable representation (a Target
Stack JSON instance). It:

1. Loads and validates the contract against schemas/contract.schema.json.
2. For each module in the contract's `infrastructure` map:
   a. Looks up the module name + version in modules/registry.json
      (version defaults to the latest non-deprecated entry when omitted).
   b. If the module is an L1 primitive: builds a stack fragment from
      the interface.json + module inputs.
   c. If the module is an L2 composition: loads the composition.json,
      expands children to stack resources, resolves wires to ref:
      expressions, and emits the fragment.
3. Merges all module fragments into a single Target Stack instance:
   - stack.name = contract.id (the short operational acronym)
   - stack.title = contract.name (the full human-readable name)
   - When the contract has one module: resource IDs are unprefixed
     (backward-compatible with existing stack consumers).
   - When the contract has multiple modules: resource IDs are prefixed
     with the module name (e.g. `microservice-vpc`) to avoid collisions,
     and all ref:/parent references are rewritten to match.

The output is a JSON instance valid against schemas/stack.schema.json,
ready for the Terraform adapter to compile.

CLI: contract_resolver.py <contract.yml> <out.json>
"""

import json
import os
import re
import sys

import yaml
import jsonschema


def _load_env(env_name, repo_root):
    """Load the environment onboarding JSON for env_name.

    Mirrors core.environment_check.load() but is self-contained so the
    resolver works both as a package import (`from core.contract_resolver
    import resolve`) and as a script (`python3 core/contract_resolver.py`).
    Emits a stderr warning when account_id is the placeholder and env != dev.
    """
    env_file = os.path.join(repo_root, "core", "environments", f"{env_name}.json")
    if not os.path.isfile(env_file):
        raise FileNotFoundError(f"no environment file for '{env_name}' at {env_file}")
    env = _load_json(env_file)
    if env.get("account_id") == "000000000000" and env_name != "dev":
        sys.stderr.write(
            f"WARNING: environment '{env_name}' has the placeholder account_id "
            f"000000000000 — replace it with the real {env_name} account id "
            f"before deploying (onboarding scaffold).\n"
        )
    return env


def _load_json(path):
    with open(path, "r") as fh:
        return json.load(fh)


def _load_yaml(path):
    with open(path, "r") as fh:
        return yaml.safe_load(fh)


_TOKEN_RE = re.compile(r"\$\{([a-zA-Z_][a-zA-Z0-9_.]*)\}")


def _lookup_dotted(context, dotted):
    """Look up a dotted path (e.g. 'env.state_backend.bucket') in context.

    context is a dict of top-level namespaces (e.g. {'env': {...}, 'contract': {...}}).
    Returns the value or raises KeyError if any segment is missing.
    """
    parts = dotted.split(".")
    cur = context
    for part in parts:
        if isinstance(cur, dict) and part in cur:
            cur = cur[part]
        else:
            raise KeyError(dotted)
    return cur


def _expand_vars(value, context):
    """Recursively expand ${env.<field>} and ${contract.<field>} tokens in value.

    Walks dicts, lists, and strings. Unknown tokens raise ValueError (fail
    loud, no silent passthrough — D-081). Dotted paths are supported
    (e.g. ${env.state_backend.bucket}). The expansion is recursive per D-087
    so nested map/list values expand too.
    """
    if isinstance(value, str):
        def _replace(match):
            token = match.group(1)
            try:
                resolved = _lookup_dotted(context, token)
            except KeyError:
                raise ValueError(f"unresolved interpolation token: ${{{token}}}")
            if isinstance(resolved, (dict, list)):
                return json.dumps(resolved)
            return str(resolved)
        return _TOKEN_RE.sub(_replace, value)
    if isinstance(value, dict):
        return {k: _expand_vars(v, context) for k, v in value.items()}
    if isinstance(value, list):
        return [_expand_vars(v, context) for v in value]
    return value


def _resolve_wire_value(wire, contract_inputs, child_outputs):
    """Resolve a wire 'from' reference to a concrete value.

    Wire 'from' can be:
    - "contract.inputs.<name>" — a contract input value
    - "<childId>.outputs.<name>" — a reference to another child's output

    Returns either a concrete value (string/number/boolean) or a
    "ref:<resourceId>.<outputName>" string for cross-child references.

    For multi-resource L1s (e.g. vpc which expands to vpc-vpc, vpc-subnet,
    vpc-routetable), the ref must point to the sub-resource that actually
    produces the output, not the child id. The child_outputs table maps
    childId -> {outputName -> resourceId} so the ref uses the correct
    resource id.
    """
    from_expr = wire["from"]
    to_expr = wire["to"]

    # If the 'from' is a contract input, use the concrete value
    if from_expr.startswith("contract.inputs."):
        input_name = from_expr[len("contract.inputs."):]
        if input_name in contract_inputs:
            return contract_inputs[input_name]
        # Check for default
        default = wire.get("default")
        if default is not None:
            return default
        return None

    # If the 'from' is a child output, emit a ref: expression
    if "." in from_expr:
        parts = from_expr.split(".", 2)
        if len(parts) >= 3 and parts[1] == "outputs":
            child_id = parts[0]
            output_name = parts[2]
            # Look up the sub-resource that produces this output.
            # child_outputs[child_id] is a dict {outputName -> resourceId}.
            # If the child is a single-resource L1, the resourceId == child_id.
            # If multi-resource, the resourceId is the expanded sub-resource id.
            child_out_map = child_outputs.get(child_id, {})
            resource_id = child_out_map.get(output_name, child_id)
            return f"ref:{resource_id}.{output_name}"

    return None


def _latest_version(registry, module_name):
    """Return the latest non-deprecated version string for a module.

    Falls back to the highest version even if all are deprecated.
    """
    versions = registry[module_name]
    non_deprecated = [(v, e) for v, e in versions.items()
                      if not e.get("deprecated", False)]
    if not non_deprecated:
        non_deprecated = list(versions.items())
    non_deprecated.sort(key=lambda x: [int(p) for p in x[0].split(".")],
                        reverse=True)
    return non_deprecated[0][0]


def _resolve_l1(module_name, version, inputs, registry, repo_root):
    """Resolve a single L1 primitive module to a stack-fragment (resources list)."""
    module_ref = f"{module_name}@{version}"

    # Load the interface
    entry = registry[module_name][version]
    iface_path = os.path.join(repo_root, entry["interface"])
    iface = _load_json(iface_path)

    # Build the resource
    resource = {
        "id": iface.get("type", module_name).split(":")[-1]
        if ":" in iface.get("type", "") else module_name,
        "type": iface["type"],
        "module": module_ref,
        "inputs": dict(inputs),
        "outputs": {
            out_name: {"type": out_spec.get("type", "string")}
            for out_name, out_spec in iface.get("outputs", {}).items()
        },
    }

    # Add NFRs if present in the interface
    nfrs = iface.get("nfrs", {})
    if nfrs:
        resource["nfrs"] = nfrs

    return {
        "kind": "l1",
        "depth": 1,
        "resources": [resource],
        "features": {},
        "outputs": {},
    }


def _resolve_l2(module_name, version, inputs, registry, repo_root):
    """Resolve a single L2 composition module to a stack-fragment.

    Returns a dict with: kind, depth, resources, features, outputs.
    The caller is responsible for merging fragments and setting stack.name/title.
    """
    # Load the composition
    entry = registry[module_name][version]
    comp_path = os.path.join(repo_root, entry["interface"])
    composition = _load_json(comp_path)

    # Track child outputs for wire resolution
    # child_outputs[childId] = {outputName -> resourceId}
    # For single-resource L1s, resourceId == childId
    # For multi-resource L1s, resourceId is the expanded sub-resource id
    child_outputs = {}
    # child_input_map[childId] = {inputName -> sub_resource_id} for multi-resource L1s
    # so a wire targeting <childId>.inputs.<name> routes to the sub-resource
    # that actually declares that input (P1-1 — desired_count -> aws:ecs:service,
    # family -> aws:ecs:task_definition).
    child_input_map = {}
    # data_source_names: set of child ids that are data sources (not modules)
    # The adapter emits `data` blocks for these instead of `module` blocks.
    data_source_names = set()
    resources = []

    # Expand children to resources
    for child in composition["children"]:
        child_id = child["id"]
        child_module = child["module"]
        child_name = child_module.split("@")[0]
        child_version = child_module.split("@")[1] if "@" in child_module else "1.0.0"

        # Load the child's interface to get type and outputs
        child_entry = registry[child_name][child_version]
        child_iface_path = os.path.join(repo_root, child_entry["interface"])
        child_iface = _load_json(child_iface_path)

        # Build the output->resourceId map for this child
        child_out_map = {}
        child_in_map = {}

        # For multi-resource L1s (like vpc), the first resource type is the
        # primary; the adapter handles expansion. Use the interface's type
        # or the first resource in the interface's resources array.
        if "resources" in child_iface and child_iface["resources"]:
            # Multi-resource L1: create one resource per sub-resource
            for sub_res in child_iface["resources"]:
                res_id = f"{child_id}-{sub_res['type'].split(':')[-1].replace('_', '-')}" if len(child_iface["resources"]) > 1 else child_id
                resource = {
                    "id": res_id,
                    "type": sub_res["type"],
                    "module": child_module,
                    "inputs": {},
                    "outputs": {
                        out: {"type": "string"}
                        for out in sub_res.get("outputs", [])
                    },
                }
                resources.append(resource)
                # Map each output to this sub-resource's id
                for out_name in sub_res.get("outputs", []):
                    child_out_map[out_name] = res_id
                # Map each declared input to this sub-resource's id (P1-1)
                for in_name in sub_res.get("inputs", []):
                    child_in_map[in_name] = res_id
        else:
            # Single-resource L1
            resource = {
                "id": child_id,
                "type": child_iface["type"],
                "module": child_module,
                "inputs": {},
                "outputs": {
                    out_name: {"type": out_spec.get("type", "string")}
                    for out_name, out_spec in child_iface.get("outputs", {}).items()
                },
            }
            resources.append(resource)
            # Map each output to the child id
            for out_name in child_iface.get("outputs", {}):
                child_out_map[out_name] = child_id

        # Also map interface-level outputs (for L1s that declare outputs at the
        # interface level rather than per-resource)
        for out_name in child_iface.get("outputs", {}):
            if out_name not in child_out_map:
                child_out_map[out_name] = child_id

        child_outputs[child_id] = child_out_map
        child_input_map[child_id] = child_in_map

    # P58: Process data_sources — pseudo-children that reference platform
    # infrastructure via terraform_remote_state. They have outputs but no
    # resources (the adapter emits `data` blocks, not `module` blocks).
    for ds in composition.get("data_sources", []):
        ds_name = ds["name"]
        data_source_names.add(ds_name)
        ds_outputs = ds.get("outputs", [])
        child_outputs[ds_name] = {out: ds_name for out in ds_outputs}

    # Resolve wires to populate inputs
    for wire in composition.get("wires", []):
        to_expr = wire["to"]
        # Parse "to": "<childId>.inputs.<inputName>"
        to_parts = to_expr.split(".")
        if len(to_parts) != 3 or to_parts[1] != "inputs":
            continue
        target_child = to_parts[0]
        input_name = to_parts[2]

        value = _resolve_wire_value(wire, inputs, child_outputs)
        if value is not None:
            # Route to the sub-resource that declares this input (P1-1).
            # child_input_map maps <childId> -> {inputName -> sub_resource_id}.
            # If the input is declared on a specific sub-resource, route there;
            # otherwise fall back to the first matching resource (legacy).
            in_map = child_input_map.get(target_child, {})
            target_res_id = in_map.get(input_name)
            if target_res_id is not None:
                for res in resources:
                    if res["id"] == target_res_id:
                        res["inputs"][input_name] = value
                        break
            else:
                for res in resources:
                    if res["id"] == target_child or res["id"].startswith(f"{target_child}-"):
                        res["inputs"][input_name] = value
                        break

    # REQ-87: Propagate deletion_protection feature flag from contract inputs
    # to all children's NFRs. When inputs.deletion_protection is false,
    # all resources get deletion_protection=false (used by decommission).
    features = {}
    deletion_protection_input = inputs.get("deletion_protection", True)
    if deletion_protection_input is not True:
        for res in resources:
            if "nfrs" not in res:
                res["nfrs"] = {}
            res["nfrs"]["deletion_protection"] = deletion_protection_input
    # Also record the feature flag on the stack object for introspection.
    if "deletion_protection" in inputs:
        features["deletion_protection"] = deletion_protection_input

    # P1-7: Process the composition's outputs[] array to build stack.outputs.
    # Each output wire: {"from": "<childId>.outputs.<name>", "to": "stack.outputs.<outName>"}
    # The child_outputs map (childId -> {outputName: resourceId}) resolves
    # the source to a resource id, which the adapter uses to emit
    # `output "<outName>" { value = aws_<type>.<resourceId>.<attr> }`.
    stack_outputs = {}
    for out_wire in composition.get("outputs", []):
        from_expr = out_wire.get("from", "")
        to_expr = out_wire.get("to", "")
        # Parse "to": "stack.outputs.<outName>"
        to_parts = to_expr.split(".")
        if len(to_parts) != 3 or to_parts[1] != "outputs":
            continue
        out_name = to_parts[2]
        # Parse "from": "<childId>.outputs.<name>"
        from_parts = from_expr.split(".")
        if len(from_parts) != 3 or from_parts[1] != "outputs":
            continue
        src_child = from_parts[0]
        src_output = from_parts[2]
        # Resolve the source resource id from child_outputs
        child_out_map = child_outputs.get(src_child, {})
        src_resource_id = child_out_map.get(src_output, src_child)
        stack_outputs[out_name] = {
            "type": "string",
            "from": src_resource_id,
            "output": src_output,
        }

    return {
        "kind": "l2",
        "depth": composition.get("depth", 1),
        "resources": resources,
        "features": features,
        "outputs": stack_outputs,
        "data_sources": list(data_source_names),
    }


def _namespace_resources(resources, module_name):
    """Prefix all resource IDs with the module name for multi-module contracts.

    Rewrites resource 'id', 'parent', and ref: expressions in inputs/outputs
    so cross-references stay consistent within the module fragment.
    """
    prefix = f"{module_name}-"
    # Build the old->new id mapping
    id_map = {res["id"]: f"{prefix}{res['id']}" for res in resources}

    def _rewrite_ref(val):
        """Recursively rewrite ref:<id>.<out> and parent:<id> strings."""
        if isinstance(val, str):
            if val.startswith("ref:"):
                # ref:<resourceId>.<outputName>
                rest = val[4:]
                if "." in rest:
                    rid, outname = rest.split(".", 1)
                    if rid in id_map:
                        return f"ref:{id_map[rid]}.{outname}"
                return val
            return val
        if isinstance(val, dict):
            return {k: _rewrite_ref(v) for k, v in val.items()}
        if isinstance(val, list):
            return [_rewrite_ref(v) for v in val]
        return val

    for res in resources:
        res["id"] = id_map[res["id"]]
        # Rewrite parent
        if "parent" in res and res["parent"] in id_map:
            res["parent"] = id_map[res["parent"]]
        # Rewrite all ref: expressions in inputs and outputs
        res["inputs"] = _rewrite_ref(res.get("inputs", {}))
        if "outputs" in res:
            res["outputs"] = _rewrite_ref(res["outputs"])

    return resources, id_map


def decommission_transform(stack_instance):
    """REQ-92: Transform a resolved stack instance for decommission.

    Sets all scalable counts to 0 and deletion_protection to false on
    every resource. Used by the decommission pipeline mode after the
    first step (disable deletion protection) has been applied.
    """
    for res in stack_instance.get("resources", []):
        if "nfrs" not in res:
            res["nfrs"] = {}
        res["nfrs"]["deletion_protection"] = False
        inputs = res.get("inputs", {})
        if "desired_count" in inputs:
            inputs["desired_count"] = 0
        if "min_capacity" in inputs:
            inputs["min_capacity"] = 0
        if "max_capacity" in inputs:
            inputs["max_capacity"] = 0
    return stack_instance


def resolve(contract_path, repo_root=None, environment_override=None):
    """Resolve a consumer contract to a Target Stack instance.

    Args:
        contract_path: Path to the contract YAML file.
        repo_root: Root of the ACDL repo (defaults to two levels up from this file).
        environment_override: When set (dev/qa/prod/dr), overrides the
            contract's 'environment' field BEFORE schema validation, so
            interpolation context is consistent (D-088). Used by
            run_platform.sh --environment.

    Returns:
        A dict representing the Target Stack instance.
    """
    if repo_root is None:
        repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

    # Load contract
    contract = _load_yaml(contract_path)

    # Apply environment override BEFORE schema validation (D-088) so the
    # schema sees the overridden value and interpolation context is consistent.
    if environment_override:
        contract["environment"] = environment_override

    # Load schemas
    contract_schema = _load_json(os.path.join(repo_root, "schemas", "contract.schema.json"))

    # Validate contract against schema
    jsonschema.validate(contract, contract_schema)

    # Interpolation (D-081): expand ${env.<field>} + ${contract.<field>}
    # tokens AFTER schema validation (the schema sees raw tokens, which are
    # valid strings) and BEFORE IR resolution (the resolver sees concrete
    # values). The env context is the loaded environment onboarding JSON.
    env_name = contract.get("environment", "dev")
    env = _load_env(env_name, repo_root)
    # Expose 'environment' as an alias for the env's 'name' field so
    # ${env.environment} resolves (the env JSON uses 'name', but contracts
    # reference the environment by ${env.environment}).
    env["environment"] = env.get("name", env_name)
    context = {"env": env, "contract": contract}

    # Expand interpolation tokens in each module's inputs
    infrastructure = contract.get("infrastructure", {})
    for module_name, module_entry in infrastructure.items():
        module_entry["inputs"] = _expand_vars(
            module_entry.get("inputs", {}), context)

    # Load registry
    registry = _load_json(os.path.join(repo_root, "modules", "registry.json"))

    # Validate every module exists in the registry, then resolve each
    module_names = list(infrastructure.keys())
    fragments = []
    for module_name in module_names:
        if module_name not in registry:
            raise ValueError(f"module '{module_name}' not found in registry")
        module_entry = infrastructure[module_name]
        # Default version to latest non-deprecated
        version = module_entry.get("version")
        if version is None:
            version = _latest_version(registry, module_name)
        elif version not in registry[module_name]:
            raise ValueError(
                f"module '{module_name}' version '{version}' not found in registry")
        module_inputs = module_entry.get("inputs", {})

        # Determine if L1 or L2
        entry = registry[module_name][version]
        interface_path = entry["interface"]
        is_l2 = "l2" in interface_path or "composition" in interface_path

        if is_l2:
            fragment = _resolve_l2(module_name, version, module_inputs,
                                   registry, repo_root)
        else:
            fragment = _resolve_l1(module_name, version, module_inputs,
                                   registry, repo_root)
        fragments.append((module_name, fragment))

    # Merge fragments into a single stack instance
    all_resources = []
    all_data_sources = []
    max_depth = 1
    any_l2 = False
    merged_features = {}
    merged_outputs = {}

    multi_module = len(fragments) > 1

    for module_name, fragment in fragments:
        if fragment["kind"] == "l2":
            any_l2 = True
        max_depth = max(max_depth, fragment["depth"])
        merged_features.update(fragment.get("features", {}))
        all_data_sources.extend(fragment.get("data_sources", []))

        if multi_module:
            # Namespace resource IDs to avoid cross-module collisions
            namespaced, id_map = _namespace_resources(
                fragment["resources"], module_name)
            # Namespace the fragment's stack outputs (from refs)
            for out_name, out_spec in fragment.get("outputs", {}).items():
                src_id = out_spec.get("from", "")
                if src_id in id_map:
                    out_spec["from"] = id_map[src_id]
                merged_outputs[f"{module_name}-{out_name}"] = out_spec
            all_resources.extend(namespaced)
        else:
            # Single module: keep IDs as-is (backward compatible)
            merged_outputs.update(fragment.get("outputs", {}))
            all_resources.extend(fragment["resources"])

    # Determine stack kind: L2 if any module is L2 or if multi-module
    if multi_module:
        kind = "l2"
    elif any_l2:
        kind = "l2"
    else:
        kind = "l1"

    stack_instance = {
        "version": "1.0.0",
        "stack": {
            "name": contract["id"],
            "kind": kind,
            "depth": max_depth,
            "environment": contract.get("environment", "dev"),
        },
        "resources": all_resources,
        "data_sources": all_data_sources,
    }

    # Add the human-readable title
    if contract.get("name"):
        stack_instance["stack"]["title"] = contract["name"]

    # Add features if any were set
    if merged_features:
        stack_instance["stack"]["features"] = merged_features

    # Add stack-level outputs
    if merged_outputs:
        stack_instance["outputs"] = merged_outputs

    # Validate against stack schema
    stack_schema = _load_json(os.path.join(repo_root, "schemas", "stack.schema.json"))
    jsonschema.validate(stack_instance, stack_schema)

    return stack_instance


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("usage: contract_resolver.py <contract.yml> <out.json> [--environment <name>]", file=sys.stderr)
        sys.exit(2)
    contract_path = sys.argv[1]
    out_path = sys.argv[2]
    env_override = None
    if "--environment" in sys.argv:
        idx = sys.argv.index("--environment")
        if idx + 1 < len(sys.argv):
            env_override = sys.argv[idx + 1]
    # Also honor the ACDL_ENVIRONMENT_OVERRIDE env var (used by run_platform.sh).
    if env_override is None and os.environ.get("ACDL_ENVIRONMENT_OVERRIDE"):
        env_override = os.environ["ACDL_ENVIRONMENT_OVERRIDE"]
    result = resolve(contract_path, environment_override=env_override)
    with open(out_path, "w") as fh:
        json.dump(result, fh, indent=2)