# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Flower command line interface `new` command."""


import io
import json
import re
import zipfile
from enum import Enum
from pathlib import Path
from string import Template
from typing import Annotated, Optional

import requests
import typer

from flwr.supercore.constant import APP_ID_PATTERN, PLATFORM_API_URL

from ..utils import (
    is_valid_project_name,
    prompt_options,
    prompt_text,
    sanitize_project_name,
)


class MlFramework(str, Enum):
    """Available frameworks."""

    PYTORCH = "PyTorch"
    TENSORFLOW = "TensorFlow"
    SKLEARN = "sklearn"
    HUGGINGFACE = "HuggingFace"
    JAX = "JAX"
    MLX = "MLX"
    NUMPY = "NumPy"
    XGBOOST = "XGBoost"
    FLOWERTUNE = "FlowerTune"
    BASELINE = "Flower Baseline"
    PYTORCH_LEGACY_API = "PyTorch (Legacy API, deprecated)"


class LlmChallengeName(str, Enum):
    """Available LLM challenges."""

    GENERALNLP = "GeneralNLP"
    FINANCE = "Finance"
    MEDICAL = "Medical"
    CODE = "Code"


class TemplateNotFound(Exception):
    """Raised when template does not exist."""


def load_template(name: str) -> str:
    """Load template from template directory and return as text."""
    tpl_dir = (Path(__file__).parent / "templates").absolute()
    tpl_file_path = tpl_dir / name

    if not tpl_file_path.is_file():
        raise TemplateNotFound(f"Template '{name}' not found")

    with open(tpl_file_path, encoding="utf-8") as tpl_file:
        return tpl_file.read()


def render_template(template: str, data: dict[str, str]) -> str:
    """Render template."""
    tpl_file = load_template(template)
    tpl = Template(tpl_file)
    if ".gitignore" not in template:
        return tpl.substitute(data)
    return tpl.template


def create_file(file_path: Path, content: str) -> None:
    """Create file including all nessecary directories and write content into file."""
    file_path.parent.mkdir(exist_ok=True)
    file_path.write_text(content, encoding="utf-8")


def render_and_create(file_path: Path, template: str, context: dict[str, str]) -> None:
    """Render template and write to file."""
    content = render_template(template, context)
    create_file(file_path, content)


def print_success_prompt(
    package_name: str, llm_challenge_str: Optional[str] = None
) -> None:
    """Print styled setup instructions for running a new Flower App after creation."""
    prompt = typer.style(
        "🎊 Flower App creation successful.\n\n"
        "To run your Flower App, first install its dependencies:\n\n",
        fg=typer.colors.GREEN,
        bold=True,
    )

    _add = "	huggingface-cli login\n" if llm_challenge_str else ""

    prompt += typer.style(
        f"	cd {package_name} && pip install -e .\n" + _add + "\n",
        fg=typer.colors.BRIGHT_CYAN,
        bold=True,
    )

    prompt += typer.style(
        "then, run the app:\n\n ",
        fg=typer.colors.GREEN,
        bold=True,
    )

    prompt += typer.style(
        "\tflwr run .\n\n",
        fg=typer.colors.BRIGHT_CYAN,
        bold=True,
    )

    prompt += typer.style(
        "💡 Check the README in your app directory to learn how to\n"
        "customize it and how to run it using the Deployment Runtime.\n",
        fg=typer.colors.GREEN,
        bold=True,
    )

    print(prompt)


# Security: prevent zip-slip
def _safe_extract_zip(zf: zipfile.ZipFile, dest_dir: Path) -> None:
    """Extract ZIP file into destination directory."""
    dest_dir = dest_dir.resolve()

    def _is_within_directory(base: Path, target: Path) -> bool:
        try:
            target.relative_to(base)
            return True
        except ValueError:
            return False

    for member in zf.infolist():
        # Skip directory placeholders;
        # ZipInfo can represent them as names ending with '/'.
        if member.is_dir():
            target_path = (dest_dir / member.filename).resolve()
            if not _is_within_directory(dest_dir, target_path):
                raise ValueError(f"Unsafe path in zip: {member.filename}")
            target_path.mkdir(parents=True, exist_ok=True)
            continue

        # Files
        target_path = (dest_dir / member.filename).resolve()
        if not _is_within_directory(dest_dir, target_path):
            raise ValueError(f"Unsafe path in zip: {member.filename}")

        # Ensure parent exists
        target_path.parent.mkdir(parents=True, exist_ok=True)

        # Extract
        with zf.open(member, "r") as src, open(target_path, "wb") as dst:
            dst.write(src.read())


def _download_zip_to_memory(presigned_url: str) -> io.BytesIO:
    """Download ZIP file from Platform API to memory."""
    try:
        r = requests.get(presigned_url, timeout=60)
        r.raise_for_status()
    except requests.RequestException as e:
        raise typer.BadParameter(f"ZIP download failed: {e}") from e

    buf = io.BytesIO(r.content)
    # Validate it's a zip
    if not zipfile.is_zipfile(buf):
        raise typer.BadParameter("Downloaded file is not a valid ZIP")
    buf.seek(0)
    return buf


def _request_download_link(identifier: str) -> str:
    """Request download link from Flower platform API."""
    url = f"{PLATFORM_API_URL}/hub/fetch-zip"
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json",
    }
    body = {
        "identifier": identifier,  # send raw string of identifier
    }

    try:
        resp = requests.post(url, headers=headers, data=json.dumps(body), timeout=20)
    except requests.RequestException as e:
        raise typer.BadParameter(f"Unable to connect to Platform API: {e}") from e

    if resp.status_code == 404:
        raise typer.BadParameter(f"'{identifier}' not found in Platform API")
    if not resp.ok:
        raise typer.BadParameter(
            f"Platform API request failed with "
            f"status {resp.status_code}. Details: {resp.text}"
        )

    data = resp.json()
    if "zip_url" not in data:
        raise typer.BadParameter("Invalid response from Platform API")
    return str(data["zip_url"])


def download_remote_app_via_api(identifier: str) -> None:
    """Download App from Platform API."""
    # Parse @user/app just to derive local dir name
    m = re.match(APP_ID_PATTERN, identifier)
    if not m:
        raise typer.BadParameter(
            "Invalid remote app ID. Expected format: '@user_name/app_name'."
        )
    app_name = m.group("app")

    project_dir = Path.cwd() / app_name
    if project_dir.exists():
        if not typer.confirm(
            typer.style(
                f"\n💬 {app_name} already exists, do you want to override it?",
                fg=typer.colors.MAGENTA,
                bold=True,
            )
        ):
            return

    print(
        typer.style(
            f"\n🔗 Requesting download link for {identifier}...",
            fg=typer.colors.GREEN,
            bold=True,
        )
    )
    presigned_url = _request_download_link(identifier)

    print(
        typer.style(
            "⬇️  Downloading ZIP into memory...",
            fg=typer.colors.GREEN,
            bold=True,
        )
    )
    zip_buf = _download_zip_to_memory(presigned_url)

    print(
        typer.style(
            f"📦 Unpacking into {project_dir}...",
            fg=typer.colors.GREEN,
            bold=True,
        )
    )
    with zipfile.ZipFile(zip_buf) as zf:
        _safe_extract_zip(zf, Path.cwd())

    print_success_prompt(app_name)


# pylint: disable=too-many-locals,too-many-branches,too-many-statements
def new(
    app_name: Annotated[
        Optional[str],
        typer.Argument(help="The name of the Flower App"),
    ] = None,
    framework: Annotated[
        Optional[MlFramework],
        typer.Option(case_sensitive=False, help="The ML framework to use"),
    ] = None,
    username: Annotated[
        Optional[str],
        typer.Option(case_sensitive=False, help="The Flower username of the author"),
    ] = None,
) -> None:
    """Create new Flower App."""
    if app_name is None:
        app_name = prompt_text("Please provide the app name")

    # Download remote app
    if app_name and app_name.startswith("@"):
        download_remote_app_via_api(app_name)
        return

    if not is_valid_project_name(app_name):
        app_name = prompt_text(
            "Please provide a name that only contains "
            "characters in {'-', a-zA-Z', '0-9'}",
            predicate=is_valid_project_name,
            default=sanitize_project_name(app_name),
        )

    # Set project directory path
    package_name = re.sub(r"[-_.]+", "-", app_name).lower()
    import_name = package_name.replace("-", "_")
    project_dir = Path.cwd() / package_name

    if project_dir.exists():
        if not typer.confirm(
            typer.style(
                f"\n💬 {app_name} already exists, do you want to override it?",
                fg=typer.colors.MAGENTA,
                bold=True,
            )
        ):
            return

    if username is None:
        username = prompt_text("Please provide your Flower username")

    if framework is not None:
        framework_str = str(framework.value)
    else:
        framework_str = prompt_options(
            "Please select ML framework by typing in the number",
            [mlf.value for mlf in MlFramework],
        )

    llm_challenge_str = None
    if framework_str == MlFramework.FLOWERTUNE:
        llm_challenge_value = prompt_options(
            "Please select LLM challenge by typing in the number",
            sorted([challenge.value for challenge in LlmChallengeName]),
        )
        llm_challenge_str = llm_challenge_value.lower()

    if framework_str == MlFramework.BASELINE:
        framework_str = "baseline"

    if framework_str == MlFramework.PYTORCH_LEGACY_API:
        framework_str = "pytorch_legacy_api"

    print(
        typer.style(
            f"\n🔨 Creating Flower App {app_name}...",
            fg=typer.colors.GREEN,
            bold=True,
        )
    )

    context = {
        "framework_str": framework_str,
        "import_name": import_name.replace("-", "_"),
        "package_name": package_name,
        "project_name": app_name,
        "username": username,
    }

    template_name = framework_str.lower()

    # List of files to render
    if llm_challenge_str:
        files = {
            ".gitignore": {"template": "app/.gitignore.tpl"},
            "pyproject.toml": {"template": f"app/pyproject.{template_name}.toml.tpl"},
            "README.md": {"template": f"app/README.{template_name}.md.tpl"},
            f"{import_name}/__init__.py": {"template": "app/code/__init__.py.tpl"},
            f"{import_name}/server_app.py": {
                "template": "app/code/flwr_tune/server_app.py.tpl"
            },
            f"{import_name}/client_app.py": {
                "template": "app/code/flwr_tune/client_app.py.tpl"
            },
            f"{import_name}/models.py": {
                "template": "app/code/flwr_tune/models.py.tpl"
            },
            f"{import_name}/dataset.py": {
                "template": "app/code/flwr_tune/dataset.py.tpl"
            },
            f"{import_name}/strategy.py": {
                "template": "app/code/flwr_tune/strategy.py.tpl"
            },
        }

        # Challenge specific context
        fraction_train = "0.2" if llm_challenge_str == "code" else "0.1"
        if llm_challenge_str == "generalnlp":
            challenge_name = "General NLP"
            num_clients = "20"
            dataset_name = "flwrlabs/alpaca-gpt4"
        elif llm_challenge_str == "finance":
            challenge_name = "Finance"
            num_clients = "50"
            dataset_name = "flwrlabs/fingpt-sentiment-train"
        elif llm_challenge_str == "medical":
            challenge_name = "Medical"
            num_clients = "20"
            dataset_name = "flwrlabs/medical-meadow-medical-flashcards"
        else:
            challenge_name = "Code"
            num_clients = "10"
            dataset_name = "flwrlabs/code-alpaca-20k"

        context["llm_challenge_str"] = llm_challenge_str
        context["fraction_train"] = fraction_train
        context["challenge_name"] = challenge_name
        context["num_clients"] = num_clients
        context["dataset_name"] = dataset_name
    else:
        files = {
            ".gitignore": {"template": "app/.gitignore.tpl"},
            "README.md": {"template": "app/README.md.tpl"},
            "pyproject.toml": {"template": f"app/pyproject.{template_name}.toml.tpl"},
            f"{import_name}/__init__.py": {"template": "app/code/__init__.py.tpl"},
            f"{import_name}/server_app.py": {
                "template": f"app/code/server.{template_name}.py.tpl"
            },
            f"{import_name}/client_app.py": {
                "template": f"app/code/client.{template_name}.py.tpl"
            },
        }

        # Depending on the framework, generate task.py file
        frameworks_with_tasks = [
            MlFramework.PYTORCH.value,
            MlFramework.JAX.value,
            MlFramework.HUGGINGFACE.value,
            MlFramework.MLX.value,
            MlFramework.TENSORFLOW.value,
            MlFramework.SKLEARN.value,
            MlFramework.NUMPY.value,
            MlFramework.XGBOOST.value,
            "pytorch_legacy_api",
        ]
        if framework_str in frameworks_with_tasks:
            files[f"{import_name}/task.py"] = {
                "template": f"app/code/task.{template_name}.py.tpl"
            }

        if framework_str == "pytorch_legacy_api":
            # Use custom __init__ that better captures name of framework
            files[f"{import_name}/__init__.py"] = {
                "template": f"app/code/__init__.{framework_str}.py.tpl"
            }

        if framework_str == "baseline":
            # Include additional files for baseline template
            for file_name in ["model", "dataset", "strategy", "utils", "__init__"]:
                files[f"{import_name}/{file_name}.py"] = {
                    "template": f"app/code/{file_name}.{template_name}.py.tpl"
                }

            # Replace README.md
            files["README.md"]["template"] = f"app/README.{template_name}.md.tpl"

            # Add LICENSE
            files["LICENSE"] = {"template": "app/LICENSE.tpl"}

    for file_path, value in files.items():
        render_and_create(
            file_path=project_dir / file_path,
            template=value["template"],
            context=context,
        )

    print_success_prompt(package_name, llm_challenge_str)
