import argparse
import shutil
import tarfile
import tempfile
import urllib.request
from dataclasses import dataclass
from pathlib import Path

import coremltools as ct
import tensorflow as tf

DEFAULT_TFHUB_URL = "https://tfhub.dev/google/yamnet/1?tf-hub-format=compressed"


def parse_args():
    script_dir = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(description="Convert YAMNet TensorFlow SavedModel to Core ML.")
    parser.add_argument(
        "--model-path",
        default=script_dir / "yamnet_model",
        type=Path,
        help="Path to a TensorFlow SavedModel directory, .keras file, or .h5 file.",
    )
    parser.add_argument(
        "--output",
        default=script_dir / "YAMNet.mlpackage",
        type=Path,
        help="Output Core ML package path.",
    )
    parser.add_argument(
        "--waveform-samples",
        default=15_600,
        type=int,
        help="Fixed waveform length for the Core ML input. YAMNet commonly uses 0.975s at 16 kHz.",
    )
    parser.add_argument(
        "--output-key",
        default=None,
        help="SavedModel output key for --conversion-mode waveform. Defaults to output_0 for TFHub YAMNet.",
    )
    parser.add_argument(
        "--conversion-mode",
        choices=["features", "waveform"],
        default="features",
        help=(
            "features converts the YAMNet classifier from 96x64 log-mel patches and avoids unsupported "
            "TensorFlow FFT ops. waveform attempts direct full SavedModel conversion."
        ),
    )
    parser.add_argument(
        "--download-tfhub",
        action="store_true",
        help="Download YAMNet from TFHub into --model-path before conversion if it is missing.",
    )
    parser.add_argument(
        "--force-download",
        action="store_true",
        help="Replace --model-path with a fresh TFHub download before conversion.",
    )
    parser.add_argument(
        "--download-only",
        action="store_true",
        help="Download YAMNet from TFHub into --model-path and exit without Core ML conversion.",
    )
    parser.add_argument(
        "--tfhub-url",
        default=DEFAULT_TFHUB_URL,
        help="TFHub compressed SavedModel URL.",
    )
    return parser.parse_args()


@dataclass(frozen=True)
class YamnetParams:
    patch_frames: int = 96
    patch_bands: int = 64
    num_classes: int = 521
    conv_padding: str = "same"
    batchnorm_center: bool = True
    batchnorm_scale: bool = False
    batchnorm_epsilon: float = 1e-4
    classifier_activation: str = "sigmoid"


YAMNET_LAYER_DEFS = [
    ("conv", [3, 3], 2, 32),
    ("separable_conv", [3, 3], 1, 64),
    ("separable_conv", [3, 3], 2, 128),
    ("separable_conv", [3, 3], 1, 128),
    ("separable_conv", [3, 3], 2, 256),
    ("separable_conv", [3, 3], 1, 256),
    ("separable_conv", [3, 3], 2, 512),
    ("separable_conv", [3, 3], 1, 512),
    ("separable_conv", [3, 3], 1, 512),
    ("separable_conv", [3, 3], 1, 512),
    ("separable_conv", [3, 3], 1, 512),
    ("separable_conv", [3, 3], 1, 512),
    ("separable_conv", [3, 3], 2, 1024),
    ("separable_conv", [3, 3], 1, 1024),
]


def download_tfhub_saved_model(tfhub_url, model_path, force=False):
    if model_path.exists() and not force:
        print(f"Using existing model: {model_path}")
        return

    if model_path.exists():
        if model_path.is_dir():
            shutil.rmtree(model_path)
        else:
            model_path.unlink()

    model_path.parent.mkdir(parents=True, exist_ok=True)
    print(f"Downloading TFHub model: {tfhub_url}")

    with tempfile.TemporaryDirectory(prefix="yamnet-tfhub-") as temp_dir:
        archive_path = Path(temp_dir) / "model.tar.gz"
        request = urllib.request.Request(tfhub_url, headers={"User-Agent": "langpipe-yamnet-coreml"})
        with urllib.request.urlopen(request) as response, archive_path.open("wb") as archive:
            shutil.copyfileobj(response, archive)

        extract_path = Path(temp_dir) / "model"
        extract_path.mkdir()
        with tarfile.open(archive_path, "r:gz") as tar:
            tar.extractall(extract_path, filter="data")

        saved_model_pb = next(extract_path.rglob("saved_model.pb"), None)
        if saved_model_pb is None:
            raise ValueError(f"TFHub archive did not contain saved_model.pb: {tfhub_url}")

        extracted_model_root = saved_model_pb.parent
        shutil.copytree(extracted_model_root, model_path)

    print(f"Saved TFHub model: {model_path}")


def batch_norm(name, params, layer_input):
    return tf.keras.layers.BatchNormalization(
        name=name,
        center=params.batchnorm_center,
        scale=params.batchnorm_scale,
        epsilon=params.batchnorm_epsilon,
    )(layer_input)


def build_yamnet_feature_model(params=YamnetParams()):
    features = tf.keras.Input(
        shape=(params.patch_frames, params.patch_bands),
        dtype=tf.float32,
        name="features",
    )
    net = tf.keras.layers.Reshape(
        (params.patch_frames, params.patch_bands, 1),
        name="features_4d",
    )(features)

    for layer_index, (layer_type, kernel, stride, filters) in enumerate(YAMNET_LAYER_DEFS, start=1):
        prefix = f"layer{layer_index}"
        if layer_type == "conv":
            net = tf.keras.layers.Conv2D(
                name=f"{prefix}_conv",
                filters=filters,
                kernel_size=kernel,
                strides=stride,
                padding=params.conv_padding,
                use_bias=False,
                activation=None,
            )(net)
            net = batch_norm(f"{prefix}_conv_bn", params, net)
            net = tf.keras.layers.ReLU(name=f"{prefix}_relu")(net)
            continue

        net = tf.keras.layers.DepthwiseConv2D(
            name=f"{prefix}_depthwise_conv",
            kernel_size=kernel,
            strides=stride,
            depth_multiplier=1,
            padding=params.conv_padding,
            use_bias=False,
            activation=None,
        )(net)
        net = batch_norm(f"{prefix}_depthwise_conv_bn", params, net)
        net = tf.keras.layers.ReLU(name=f"{prefix}_depthwise_relu")(net)
        net = tf.keras.layers.Conv2D(
            name=f"{prefix}_pointwise_conv",
            filters=filters,
            kernel_size=(1, 1),
            strides=1,
            padding=params.conv_padding,
            use_bias=False,
            activation=None,
        )(net)
        net = batch_norm(f"{prefix}_pointwise_conv_bn", params, net)
        net = tf.keras.layers.ReLU(name=f"{prefix}_pointwise_relu")(net)

    embeddings = tf.keras.layers.GlobalAveragePooling2D(name="embeddings")(net)
    logits = tf.keras.layers.Dense(units=params.num_classes, use_bias=True, name="dense")(embeddings)
    class_scores = tf.keras.layers.Activation(
        activation=params.classifier_activation,
        name="class_scores",
    )(logits)
    return tf.keras.Model(name="yamnet_features", inputs=features, outputs=class_scores)


def load_feature_model_weights_from_saved_model(model, model_path):
    loaded = tf.saved_model.load(str(model_path))
    if not hasattr(loaded, "_yamnet"):
        raise ValueError("SavedModel does not expose the expected TFHub YAMNet _yamnet object.")

    source_weights = list(loaded._yamnet.variables)
    target_weights = model.weights
    if len(source_weights) != len(target_weights):
        raise ValueError(f"Weight count mismatch: source={len(source_weights)}, target={len(target_weights)}")

    for index, (target, source) in enumerate(zip(target_weights, source_weights)):
        if tuple(target.shape) != tuple(source.shape):
            raise ValueError(
                f"Weight shape mismatch at {index}: target {target.name} {target.shape}, "
                f"source {source.name} {source.shape}"
            )

    model.set_weights([weight.numpy() for weight in source_weights])


def describe_signature(signature):
    _, keyword_specs = signature.structured_input_signature
    print("SavedModel inputs:")
    for name, spec in keyword_specs.items():
        print(f"  {name}: shape={spec.shape}, dtype={spec.dtype.name}")

    print("SavedModel outputs:")
    for name, spec in signature.structured_outputs.items():
        print(f"  {name}: shape={spec.shape}, dtype={spec.dtype.name}")


def select_output_key(signature, requested_key):
    outputs = signature.structured_outputs
    if requested_key:
        if requested_key not in outputs:
            raise ValueError(f"--output-key {requested_key!r} not found. Available keys: {list(outputs)}")
        return requested_key

    if "output_0" in outputs:
        return "output_0"

    for key in outputs:
        if "score" in key.lower() or "class" in key.lower():
            return key

    if not outputs:
        raise ValueError("SavedModel serving_default has no outputs.")
    return next(iter(outputs))


class SavedModelClassScores(tf.Module):
    def __init__(self, signature, input_key, output_key):
        super().__init__()
        self.signature = signature
        self.input_key = input_key
        self.output_key = output_key

    @tf.function
    def __call__(self, waveform):
        outputs = self.signature(**{self.input_key: waveform})
        return {"class_scores": outputs[self.output_key]}


def convert_saved_model(model_path, output_path, waveform_samples, output_key):
    loaded = tf.saved_model.load(str(model_path))
    if "serving_default" not in loaded.signatures:
        raise ValueError(f"SavedModel has no serving_default signature. Available: {list(loaded.signatures)}")

    signature = loaded.signatures["serving_default"]
    describe_signature(signature)

    _, keyword_specs = signature.structured_input_signature
    if len(keyword_specs) != 1:
        raise ValueError(
            "Expected one SavedModel input. Pass a wrapper model if this SavedModel has "
            f"{len(keyword_specs)} inputs: {list(keyword_specs)}"
        )

    input_key = next(iter(keyword_specs))
    selected_output = select_output_key(signature, output_key)
    print(f"Converting input {input_key!r} -> output {selected_output!r} as 'class_scores'")

    wrapper = SavedModelClassScores(signature, input_key, selected_output)
    concrete = wrapper.__call__.get_concrete_function(
        tf.TensorSpec([waveform_samples], tf.float32, name="waveform")
    )

    return ct.convert(
        [concrete],
        source="tensorflow",
        inputs=[ct.TensorType(shape=(waveform_samples,), name="waveform")],
        convert_to="mlprogram",
        minimum_deployment_target=ct.target.macOS13,
    )


def convert_keras_model(model_path, output_path):
    model = tf.keras.models.load_model(str(model_path))
    return ct.convert(
        model,
        inputs=[ct.TensorType(shape=(1, 64, 96, 1), name="features")],
        outputs=[ct.TensorType(name="class_scores")],
        convert_to="mlprogram",
        minimum_deployment_target=ct.target.macOS13,
    )


def convert_feature_model(model_path):
    model = build_yamnet_feature_model()
    model(tf.zeros((1, 96, 64), dtype=tf.float32))
    load_feature_model_weights_from_saved_model(model, model_path)
    return ct.convert(
        model,
        source="tensorflow",
        inputs=[ct.TensorType(shape=(1, 96, 64), name="features")],
        convert_to="mlprogram",
        minimum_deployment_target=ct.target.macOS13,
    )


def main():
    args = parse_args()
    if args.download_tfhub or args.force_download or args.download_only:
        download_tfhub_saved_model(args.tfhub_url, args.model_path, force=args.force_download)
    if args.download_only:
        return

    if not args.model_path.exists():
        raise FileNotFoundError(
            f"Model path does not exist: {args.model_path}. "
            "Pass --download-tfhub to download YAMNet from TFHub."
        )

    if args.conversion_mode == "features":
        mlmodel = convert_feature_model(args.model_path)
    elif args.model_path.suffix in {".keras", ".h5", ".hdf5"}:
        mlmodel = convert_keras_model(args.model_path, args.output)
    else:
        mlmodel = convert_saved_model(
            args.model_path,
            args.output,
            args.waveform_samples,
            args.output_key,
        )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    mlmodel.save(str(args.output))
    print(f"Saved Core ML model: {args.output}")


if __name__ == "__main__":
    main()