# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
"""Progressive rollout setup utilities for Airbyte connectors.

This module provides pure Python functionality to configure a connector's
metadata.yaml for progressive rollout of a release candidate (RC) version.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import semver
import yaml

from airbyte_ops_mcp.airbyte_repo.bump_version import (
    InvalidVersionError,
    get_connector_path,
    get_current_version,
)
from airbyte_ops_mcp.airbyte_repo.list_connectors import METADATA_FILE_NAME


@dataclass
class ProgressiveRolloutSetupResult:
    """Result of a progressive rollout setup operation."""

    connector: str
    previous_version: str
    rc_version: str
    pin_version: str
    files_modified: list[str]
    dry_run: bool


@dataclass
class ProgressiveRolloutCleanupResult:
    """Result of a progressive rollout cleanup operation."""

    connector: str
    previous_version: str
    new_version: str
    files_modified: list[str]
    dry_run: bool
    changes_made: list[str]


def calculate_rc_version(
    current_version: str,
    rc_version: str | None = None,
) -> str:
    """Calculate the RC version for progressive rollout.

    If rc_version is provided, validates and returns it.
    Otherwise, bumps the minor version and adds -rc.1 suffix.

    Args:
        current_version: Current version string (e.g., "0.0.46")
        rc_version: Explicit RC version (e.g., "0.1.0-rc.1")

    Returns:
        RC version string

    Raises:
        InvalidVersionError: If version format is invalid
    """
    if rc_version is not None:
        if not rc_version.endswith("-rc.1") and "-rc." not in rc_version:
            raise InvalidVersionError(
                f"RC version must contain '-rc.' suffix, got: {rc_version}"
            )
        return rc_version

    try:
        version = semver.Version.parse(current_version)
    except ValueError as e:
        raise InvalidVersionError(
            f"Cannot parse current version: {current_version}"
        ) from e

    bumped = version.bump_minor()
    return f"{bumped}-rc.1"


def setup_progressive_rollout(
    repo_path: str | Path,
    connector_name: str,
    rc_version: str | None = None,
    pin_version: str | None = None,
    dry_run: bool = False,
) -> ProgressiveRolloutSetupResult:
    """Configure a connector's metadata.yaml for progressive rollout.

    This function modifies metadata.yaml to:
    1. Set dockerImageTag to the RC version
    2. Add registryOverrides.cloud.dockerImageTag to pin Cloud users to current version
    3. Add registryOverrides.oss.dockerImageTag to pin OSS users to current version
    4. Add releases.rolloutConfiguration.enableProgressiveRollout: true

    Args:
        repo_path: Path to the Airbyte monorepo
        connector_name: Technical name of the connector (e.g., "source-github")
        rc_version: Explicit RC version (default: bumps minor and adds -rc.1)
        pin_version: Version to pin Cloud/OSS users to (default: current version)
        dry_run: If True, don't actually modify files

    Returns:
        ProgressiveRolloutSetupResult with details of the operation

    Raises:
        ConnectorNotFoundError: If connector doesn't exist
        VersionNotFoundError: If current version cannot be found
        InvalidVersionError: If version format is invalid
    """
    repo_path = Path(repo_path)
    connector_path = get_connector_path(repo_path, connector_name)
    metadata_file = connector_path / METADATA_FILE_NAME

    current_version = get_current_version(connector_path)
    calculated_rc_version = calculate_rc_version(current_version, rc_version)
    calculated_pin_version = pin_version if pin_version else current_version

    with open(metadata_file) as f:
        metadata = yaml.safe_load(f)

    data = metadata.get("data", {})

    data["dockerImageTag"] = calculated_rc_version

    if "registryOverrides" not in data:
        data["registryOverrides"] = {}

    if "cloud" not in data["registryOverrides"]:
        data["registryOverrides"]["cloud"] = {}
    data["registryOverrides"]["cloud"]["dockerImageTag"] = calculated_pin_version

    if "oss" not in data["registryOverrides"]:
        data["registryOverrides"]["oss"] = {}
    data["registryOverrides"]["oss"]["dockerImageTag"] = calculated_pin_version

    if "releases" not in data:
        data["releases"] = {}
    if "rolloutConfiguration" not in data["releases"]:
        data["releases"]["rolloutConfiguration"] = {}
    data["releases"]["rolloutConfiguration"]["enableProgressiveRollout"] = True

    metadata["data"] = data

    files_modified: list[str] = []

    if not dry_run:
        with open(metadata_file, "w") as f:
            yaml.dump(metadata, f, default_flow_style=False, sort_keys=False)
        files_modified.append(str(metadata_file.relative_to(repo_path)))
    else:
        files_modified.append(str(metadata_file.relative_to(repo_path)))

    return ProgressiveRolloutSetupResult(
        connector=connector_name,
        previous_version=current_version,
        rc_version=calculated_rc_version,
        pin_version=calculated_pin_version,
        files_modified=files_modified,
        dry_run=dry_run,
    )


def extract_ga_version(rc_version: str) -> str:
    """Extract the GA version from an RC version string.

    Args:
        rc_version: RC version string (e.g., "1.0.0-rc.1")

    Returns:
        GA version string (e.g., "1.0.0")

    Raises:
        InvalidVersionError: If version format is invalid
    """
    if "-rc." in rc_version:
        return rc_version.split("-rc.")[0]
    if "-preview." in rc_version:
        return rc_version.split("-preview.")[0]
    return rc_version


def cleanup_progressive_rollout(
    repo_path: str | Path,
    connector_name: str,
    ga_version: str | None = None,
    dry_run: bool = False,
) -> ProgressiveRolloutCleanupResult:
    """Clean up a connector's metadata.yaml after progressive rollout completion.

    This function reverses the changes made by setup_progressive_rollout:
    1. Updates dockerImageTag to the GA version (removes -rc.X suffix)
    2. Removes registryOverrides.cloud.dockerImageTag (if it was a pin)
    3. Removes registryOverrides.oss.dockerImageTag (if it was a pin)
    4. Sets releases.rolloutConfiguration.enableProgressiveRollout to false

    Args:
        repo_path: Path to the Airbyte monorepo
        connector_name: Technical name of the connector (e.g., "source-github")
        ga_version: Explicit GA version (default: extracts from current RC version)
        dry_run: If True, don't actually modify files

    Returns:
        ProgressiveRolloutCleanupResult with details of the operation

    Raises:
        ConnectorNotFoundError: If connector doesn't exist
        VersionNotFoundError: If current version cannot be found
        InvalidVersionError: If version format is invalid
    """
    repo_path = Path(repo_path)
    connector_path = get_connector_path(repo_path, connector_name)
    metadata_file = connector_path / METADATA_FILE_NAME

    current_version = get_current_version(connector_path)

    new_version = ga_version or extract_ga_version(current_version)

    with open(metadata_file) as f:
        metadata = yaml.safe_load(f)

    data = metadata.get("data", {})
    changes_made: list[str] = []

    if data.get("dockerImageTag") != new_version:
        data["dockerImageTag"] = new_version
        changes_made.append(f"Updated dockerImageTag to {new_version}")

    registry_overrides = data.get("registryOverrides", {})
    cloud_overrides = registry_overrides.get("cloud", {})
    oss_overrides = registry_overrides.get("oss", {})

    if "dockerImageTag" in cloud_overrides:
        del cloud_overrides["dockerImageTag"]
        changes_made.append("Removed registryOverrides.cloud.dockerImageTag")
        if not cloud_overrides:
            del registry_overrides["cloud"]
            changes_made.append("Removed empty registryOverrides.cloud section")

    if "dockerImageTag" in oss_overrides:
        del oss_overrides["dockerImageTag"]
        changes_made.append("Removed registryOverrides.oss.dockerImageTag")
        if not oss_overrides:
            del registry_overrides["oss"]
            changes_made.append("Removed empty registryOverrides.oss section")

    if not registry_overrides and "registryOverrides" in data:
        del data["registryOverrides"]
        changes_made.append("Removed empty registryOverrides section")

    releases = data.get("releases", {})
    rollout_config = releases.get("rolloutConfiguration", {})
    if rollout_config.get("enableProgressiveRollout"):
        rollout_config["enableProgressiveRollout"] = False
        changes_made.append("Set enableProgressiveRollout to false")

    metadata["data"] = data

    files_modified: list[str] = []

    if not dry_run:
        with open(metadata_file, "w") as f:
            yaml.dump(metadata, f, default_flow_style=False, sort_keys=False)
        files_modified.append(str(metadata_file.relative_to(repo_path)))
    else:
        files_modified.append(str(metadata_file.relative_to(repo_path)))

    return ProgressiveRolloutCleanupResult(
        connector=connector_name,
        previous_version=current_version,
        new_version=new_version,
        files_modified=files_modified,
        dry_run=dry_run,
        changes_made=changes_made,
    )
