#! /usr/bin/env python3
import json
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional

import arguably
import toml


def usage():
    """Print usage information and exit."""
    print("Usage: release {prod|dev} version")
    sys.exit(1)


def validate_version(environment: str, version: str) -> None:
    """Validate version format based on environment."""
    if environment == "prod":
        # For prod: semantic versioning (x.y.z)
        if not re.match(r"^\d+\.\d+\.\d+$", version):
            print("Error: For 'prod', the version must be in the format 'number.number.number'.")
            sys.exit(1)
    else:  # dev
        # For dev: Python standards (.dev, .a, .b, .rc postfix)
        if not re.match(r"^\d+\.\d+\.\d+(\.(dev|a|b|rc)\d*)?$", version):
            print(
                "Error: For 'dev', the version must follow Python standards (e.g., '1.2.3.dev0', '1.2.3.a1', '1.2.3.b2')."
            )
            sys.exit(1)


def run_command(command: str, cwd: Optional[str] = None) -> None:
    """Run a shell command with error handling."""
    try:
        subprocess.run(command, shell=True, check=True, cwd=cwd)
    except subprocess.CalledProcessError as e:
        print(f"Error: Command '{command}' failed with exit code {e.returncode}")
        sys.exit(1)


def build_package() -> None:
    """Build the Python package."""
    print("Building package...")

    # Clear dist folder if it exists
    dist_path = Path("dist")
    if dist_path.exists():
        shutil.rmtree(dist_path)
        print("Cleared dist/ folder")

    # Run build command
    run_command("python -m build")
    print("Package built successfully")


def update_pyproject_version(version: str) -> None:
    """Update version in pyproject.toml."""
    pyproject_path = Path("pyproject.toml")

    if not pyproject_path.exists():
        print("Error: pyproject.toml not found")
        sys.exit(1)

    # Read pyproject.toml
    with open(pyproject_path) as f:
        pyproject_data = toml.load(f)

    # Update version
    if "project" in pyproject_data:
        pyproject_data["project"]["version"] = version
        print(f"Updated version to {version}")
    else:
        print("Error: Could not find [project] section in pyproject.toml")
        sys.exit(1)

    # Write back to file
    with open(pyproject_path, "w") as f:
        toml.dump(pyproject_data, f)


def publish_package(environment: str) -> None:
    """Publish package to PyPI or TestPyPI."""
    print(f"Publishing to {'TestPyPI' if environment == 'dev' else 'PyPI'}...")

    if environment == "dev":
        # Publish to TestPyPI
        run_command("python -m twine upload --repository testpypi dist/*")
    else:
        # Publish to PyPI
        run_command("python -m twine upload dist/*")

    print("Package published successfully")


def update_cookbook(project_root: Path, version: str) -> None:
    """Update Python cookbook with new runtime version."""
    python_recipes_path = project_root / "apps/web/packagerWorkerAssets/cookbook/python-examples"

    if not python_recipes_path.exists():
        print(f"Warning: Cookbook path does not exist: {python_recipes_path}")
        return

    # Find all recipe directories
    recipes = [d for d in python_recipes_path.iterdir() if d.is_dir()]

    for recipe in recipes:
        pyproject_file = recipe / "pyproject.toml"
        if pyproject_file.exists():
            print(f"Updating {recipe}")
            update_recipe_pyproject(pyproject_file, version)

    # Rebuild templates from updated cookbook
    web_dir = project_root / "apps/web"
    print("Generating templates from updated cookbook...")
    run_command("yarn cookbook:generate", cwd=str(web_dir))

    # Update web packageVersions.json
    package_versions_path = web_dir / "lib/ProjectFiles/packageVersions.json"

    try:
        package_version_content = package_versions_path.read_text()
        package_versions = json.loads(package_version_content)
        package_versions["python"]["runtime"]["name"] = "intuned-runtime"
        package_versions["python"]["runtime"]["version"] = version
        package_versions_path.write_text(json.dumps(package_versions, indent=2))
    except Exception as e:
        print(
            "Warning: Could not update packageVersions.json:",
            e,
            "\n\nPlease update it manually to avoid inconsistencies.",
        )

    # Update ASM packageVersions.json
    asm_package_versions_path = project_root / "apps/authoring-server/assets/packageVersions.json"

    try:
        asm_package_versions = json.loads(asm_package_versions_path.read_text())
        asm_package_versions["python"]["runtime"]["name"] = "intuned-runtime"
        asm_package_versions["python"]["runtime"]["version"] = version
        asm_package_versions_path.write_text(json.dumps(asm_package_versions, indent=2))
    except Exception as e:
        print(
            "Warning: Could not update ASM packageVersions.json:",
            e,
            "\n\nPlease update it manually to avoid inconsistencies.",
        )


def update_recipe_pyproject(pyproject_file: Path, version: str) -> None:
    """Update pyproject.toml in recipe with exact version."""

    run_command(f"uv add --no-sync --no-cache intuned-runtime=={version}", cwd=str(pyproject_file.parent))


@arguably.command  # type: ignore
def release(environment: str, version: str):
    """
    Release a new version of the intuned-runtime package.

    Args:
        environment (str): 'prod' for production release, 'dev' for development release.
        version (str): Version string to set for the release.
    """

    if environment not in ["prod", "dev"]:
        print("Error: The first argument must be either 'prod' or 'dev'.")
        usage()

    validate_version(environment, version)

    print(f"Environment: {environment}")
    print(f"Version: {version}")

    # Get project root directory
    script_dir = Path(__file__).parent
    project_root = script_dir.parent.parent.parent

    try:
        # Update version in pyproject.toml
        update_pyproject_version(version)

        # Sync dependencies to update uv.lock after version changes
        run_command("uv sync")

        # Build the package
        build_package()

        # Publish the package
        publish_package(environment)

        if environment == "prod":
            # Wait for package to be available
            print("Waiting 20 seconds for package to be available...")
            time.sleep(20)

            # Update cookbook 
            update_cookbook(project_root, version)

            print(f"Successfully updated intuned-runtime to version {version}")
            print("To keep cookbook in sync, create a new branch and open a PR on the cookbook repo")
        else:
            print(f"Successfully updated intuned-runtime to version {version}")

            print(f"""Use "intuned-runtime==={version}" to install this development version from TestPyPI.

Also add the following to your `uv` configuration:
[[tool.uv.index]]
name = "testpypi"
url = "https://test.pypi.org/simple/"
explicit = true

[tool.uv.sources]
intuned-runtime = {{ index = "testpypi" }}
""")

    except Exception as error:
        print(f"Error: {error}")
        sys.exit(1)


if __name__ == "__main__":
    arguably.run(name="release")
