"""
Cube MCP Server

Tools:
- cube_login: Authenticate to a Cube via Teleport
- cube_status: Get Cube node status via kubectl describe node
- cube_list: List available Cubes via Teleport
- acr_get_token: Get OAuth2 access token from Apollo Container Registry
- acr_login: Login to Apollo Container Registry (Docker + Helm)
- build_and_publish_to_apollo: Build Docker image, package Helm chart, and publish to Apollo
"""

import subprocess
import os
import re
import json
import glob
from mcp.server import FastMCP

# EdgescaleAI Teleport proxy URL
DEFAULT_PROXY = "edgescaleai.teleport.sh"

# Apollo Container Registry (ACR) configuration
ACR_REGISTRY = "edgescaleai.palantirapollo.com"
ACR_DOCKER_PREFIX = f"{ACR_REGISTRY}/com.edgescaleai-cube"
ACR_HELM_PREFIX = f"oci://{ACR_REGISTRY}/charts/com.edgescaleai-cube"
ACR_TOKEN_URL = f"https://{ACR_REGISTRY}/multipass/api/oauth2/token"
APOLLO_URL = f"https://{ACR_REGISTRY}"
APOLLO_GROUP_ID = "com.edgescaleai-cube"

mcp = FastMCP("cube-mcp")


# ============================================================================
# Helper Functions for ACR/Helm Operations
# ============================================================================

def _get_apollo_credentials():
    """Get Apollo client credentials from environment variables."""
    client_id = os.environ.get("APOLLO_CLIENT")
    client_secret = os.environ.get("APOLLO_SECRET")
    return client_id, client_secret


def _get_acr_token() -> tuple[str, str]:
    """
    Get OAuth2 token from Apollo. Returns (token, error_message).
    """
    client_id, client_secret = _get_apollo_credentials()

    if not client_id or not client_secret:
        return None, """Apollo credentials not found.

Set these environment variables in ~/.bashrc or ~/.zshrc:

  export APOLLO_CLIENT="<your-client-id>"
  export APOLLO_SECRET="<your-client-secret>"

Then restart your terminal or run: source ~/.bashrc"""

    try:
        result = subprocess.run(
            [
                "curl", "-s", ACR_TOKEN_URL,
                "-H", "content-type: application/x-www-form-urlencoded",
                "--data-urlencode", "grant_type=client_credentials",
                "--data-urlencode", f"client_id={client_id}",
                "--data-urlencode", f"client_secret={client_secret}"
            ],
            capture_output=True,
            text=True,
            timeout=30
        )

        if result.returncode != 0:
            return None, f"Failed to get token: {result.stderr}"

        response = json.loads(result.stdout)
        token = response.get("access_token")

        if not token:
            return None, f"No access_token in response: {result.stdout}"

        return token, None

    except subprocess.TimeoutExpired:
        return None, "Token request timed out"
    except json.JSONDecodeError as e:
        return None, f"Invalid JSON response: {e}"
    except FileNotFoundError:
        return None, "curl not found. Please install curl."


def _find_chart_yaml(path: str = None) -> tuple[str, str]:
    """
    Find Chart.yaml. Returns (chart_path, error_message).

    Search order:
    1. If path specified, use it directly
    2. ./Chart.yaml
    3. ./charts/*/Chart.yaml
    """
    if path:
        if os.path.isfile(path):
            return path, None
        chart_yaml = os.path.join(path, "Chart.yaml")
        if os.path.isfile(chart_yaml):
            return chart_yaml, None
        return None, f"Chart.yaml not found at: {path}"

    # Check current directory
    if os.path.isfile("Chart.yaml"):
        return "Chart.yaml", None

    # Check charts subdirectory
    charts = glob.glob("charts/*/Chart.yaml")
    if len(charts) == 1:
        return charts[0], None
    elif len(charts) > 1:
        chart_names = [os.path.dirname(c) for c in charts]
        return None, f"Multiple charts found: {', '.join(chart_names)}. Please specify which one."

    return None, "No Chart.yaml found. Run from a chart directory or specify the path."


def _parse_chart_yaml(chart_path: str) -> tuple[dict, str]:
    """
    Parse Chart.yaml and return dict with name, version, appVersion.
    Returns (chart_info, error_message).
    """
    try:
        with open(chart_path, 'r') as f:
            content = f.read()

        # Simple YAML parsing for the fields we need
        name_match = re.search(r'^name:\s*(.+)$', content, re.MULTILINE)
        version_match = re.search(r'^version:\s*(.+)$', content, re.MULTILINE)
        app_version_match = re.search(r'^appVersion:\s*["\']?([^"\']+)["\']?$', content, re.MULTILINE)

        if not name_match or not version_match:
            return None, f"Could not parse name/version from {chart_path}"

        return {
            "name": name_match.group(1).strip(),
            "version": version_match.group(1).strip(),
            "appVersion": app_version_match.group(1).strip() if app_version_match else None,
            "path": chart_path,
            "dir": os.path.dirname(chart_path) or "."
        }, None

    except Exception as e:
        return None, f"Error reading {chart_path}: {e}"


def _bump_version(version: str, bump_type: str = "patch") -> str:
    """Bump a semver version string."""
    parts = version.split(".")
    if len(parts) != 3:
        return version  # Can't bump non-semver

    major, minor, patch = int(parts[0]), int(parts[1]), int(parts[2])

    if bump_type == "major":
        return f"{major + 1}.0.0"
    elif bump_type == "minor":
        return f"{major}.{minor + 1}.0"
    else:  # patch
        return f"{major}.{minor}.{patch + 1}"


def _update_chart_version(chart_path: str, new_version: str, new_app_version: str = None) -> tuple[bool, str]:
    """Update version (and optionally appVersion) in Chart.yaml."""
    try:
        with open(chart_path, 'r') as f:
            content = f.read()

        # Update version
        content = re.sub(
            r'^(version:\s*).+$',
            f'\\g<1>{new_version}',
            content,
            flags=re.MULTILINE
        )

        # Update appVersion if specified
        if new_app_version:
            content = re.sub(
                r'^(appVersion:\s*)["\']?[^"\']+["\']?$',
                f'\\g<1>"{new_app_version}"',
                content,
                flags=re.MULTILINE
            )

        with open(chart_path, 'w') as f:
            f.write(content)

        return True, None

    except Exception as e:
        return False, f"Error updating {chart_path}: {e}"


@mcp.tool()
def cube_login(cluster: str, proxy: str = DEFAULT_PROXY) -> str:
    """
    Login to a Cube cluster via Teleport (proxy: edgescaleai.teleport.sh).
    This will authenticate you via SSO and set up kubectl access.

    Args:
        cluster: Cube cluster name to connect to (e.g., staging-int, staging-pre-prod, ironmountain-v4)
        proxy: Teleport proxy URL. Defaults to edgescaleai.teleport.sh
    """
    try:
        # Step 1: Login to Teleport (may open browser for SSO)
        login_result = subprocess.run(
            ["tsh", "login", f"--proxy={proxy}"],
            capture_output=True,
            text=True,
            timeout=120  # SSO flow can take time
        )

        if login_result.returncode != 0:
            return f"Teleport login failed:\n{login_result.stderr}\n\nTry running manually: tsh login --proxy={proxy}"

        # Step 2: Connect to the Cube cluster
        kube_result = subprocess.run(
            ["tsh", "kube", "login", cluster],
            capture_output=True,
            text=True,
            timeout=30
        )

        if kube_result.returncode != 0:
            return f"Failed to connect to Cube '{cluster}':\n{kube_result.stderr}\n\nUse cube_list to see available Cubes."

        return f"Successfully connected to Cube: {cluster}\n\nYou can now use:\n- cube_status: Check node health\n- kubectl commands to manage the Cube"

    except subprocess.TimeoutExpired:
        return f"Login timed out. This usually means the SSO flow wasn't completed in the browser.\n\nTry running manually: tsh login --proxy={proxy}"
    except FileNotFoundError:
        return "Teleport client (tsh) not found.\n\nInstall it from: https://goteleport.com/download/"


@mcp.tool()
def cube_list() -> str:
    """
    List all available Cube clusters accessible via Teleport.
    Requires being logged into Teleport first.
    """
    try:
        # Check if logged in first
        status_result = subprocess.run(
            ["tsh", "status"],
            capture_output=True,
            text=True,
            timeout=10
        )

        if status_result.returncode != 0 or "EXPIRED" in status_result.stdout:
            return f"Not logged into Teleport or session expired.\n\nRun cube_login first to authenticate, or manually:\n  tsh login --proxy={DEFAULT_PROXY}"

        # List available Kubernetes clusters
        list_result = subprocess.run(
            ["tsh", "kube", "ls"],
            capture_output=True,
            text=True,
            timeout=30
        )

        if list_result.returncode != 0:
            return f"Failed to list Cubes:\n{list_result.stderr}"

        return f"Available Cubes:\n\n{list_result.stdout}\n\nUse cube_login with a cluster name to connect."

    except subprocess.TimeoutExpired:
        return "Command timed out. Check your network connection."
    except FileNotFoundError:
        return "Teleport client (tsh) not found.\n\nInstall it from: https://goteleport.com/download/"


@mcp.tool()
def cube_status() -> str:
    """
    Get the status of the connected Cube node.
    Shows node conditions, capacity, and resource usage.
    """
    try:
        # First, get the node name
        get_nodes = subprocess.run(
            ["kubectl", "--insecure-skip-tls-verify", "get", "nodes", "-o", "jsonpath={.items[0].metadata.name}"],
            capture_output=True,
            text=True,
            timeout=30
        )

        if get_nodes.returncode != 0:
            return f"Failed to get node. Are you connected to a Cube?\n\nError: {get_nodes.stderr}\n\nUse cube_login first to connect."

        node_name = get_nodes.stdout.strip()
        if not node_name:
            return "No nodes found. The Cube may not be running or you may not be connected."

        # Get detailed node status
        describe_result = subprocess.run(
            ["kubectl", "--insecure-skip-tls-verify", "describe", "node", node_name],
            capture_output=True,
            text=True,
            timeout=30
        )

        if describe_result.returncode != 0:
            return f"Failed to describe node:\n{describe_result.stderr}"

        return f"Cube Node Status ({node_name}):\n\n{describe_result.stdout}"

    except subprocess.TimeoutExpired:
        return "Command timed out. The Cube may be unresponsive."
    except FileNotFoundError:
        return "kubectl not found. Ensure it's installed and in your PATH."


# ============================================================================
# ACR (Apollo Container Registry) Tools
# ============================================================================

@mcp.tool()
def acr_get_token() -> str:
    """
    Get an OAuth2 access token from Apollo Container Registry.

    Requires APOLLO_CLIENT and APOLLO_SECRET environment variables.
    Useful for debugging or manual operations.
    """
    token, error = _get_acr_token()
    if error:
        return error

    return f"Successfully retrieved ACR token.\n\nToken (first 20 chars): {token[:20]}...\n\nUse acr_login to authenticate Docker and Helm."


@mcp.tool()
def acr_login() -> str:
    """
    Login to Apollo Container Registry for both Docker and Helm.

    Requires APOLLO_CLIENT and APOLLO_SECRET environment variables.
    This authenticates you to push Docker images and Helm charts to ACR.
    """
    token, error = _get_acr_token()
    if error:
        return error

    results = []

    # Docker login
    try:
        docker_result = subprocess.run(
            ["docker", "login", ACR_REGISTRY, "--username", "apollo", "--password-stdin"],
            input=token,
            capture_output=True,
            text=True,
            timeout=30
        )
        if docker_result.returncode == 0:
            results.append(f"✓ Docker logged in to {ACR_REGISTRY}")
        else:
            results.append(f"✗ Docker login failed: {docker_result.stderr}")
    except FileNotFoundError:
        results.append("✗ Docker not found (skipping)")
    except Exception as e:
        results.append(f"✗ Docker login error: {e}")

    # Helm registry login
    try:
        helm_result = subprocess.run(
            ["helm", "registry", "login", ACR_REGISTRY, "--username", "apollo", "--password-stdin"],
            input=token,
            capture_output=True,
            text=True,
            timeout=30
        )
        if helm_result.returncode == 0:
            results.append(f"✓ Helm logged in to {ACR_REGISTRY}")
        else:
            results.append(f"✗ Helm login failed: {helm_result.stderr}")
    except FileNotFoundError:
        results.append("✗ Helm not found")
    except Exception as e:
        results.append(f"✗ Helm login error: {e}")

    return "ACR Login Results:\n\n" + "\n".join(results)


@mcp.tool()
def build_and_publish_to_apollo(
    chart_path: str = None,
    version: str = None,
    bump: str = "patch",
    app_version: str = None,
    skip_docker: bool = False,
    dry_run: bool = False
) -> str:
    """
    Build and publish an app to Apollo Container Registry.

    This builds your app and publishes it to Apollo, making it available
    for deployment to Cube environments. It does NOT deploy to a Cube -
    that's a separate step done from Apollo.

    Pipeline:
    1. Login to ACR (Docker + Helm registries)
    2. Build and push Docker image (if Dockerfile found, multi-arch)
    3. Bump Chart.yaml version
    4. Package Helm chart
    5. Push Helm chart to ACR (OCI registry)
    6. Publish manifest to Apollo (makes it available for deployment)

    Args:
        chart_path: Path to chart directory or Chart.yaml (auto-detected if not specified)
        version: Exact version to use (e.g., "1.0.0"). If not specified, bumps current version.
        bump: Version bump type: "patch" (default), "minor", or "major". Ignored if version is specified.
        app_version: New appVersion for Docker image tag. If not specified, uses chart version.
        skip_docker: Skip Docker build/push even if Dockerfile exists.
        dry_run: Show what would be done without actually doing it.
    """
    steps = []
    errors = []

    # Step 1: Find and parse Chart.yaml
    chart_yaml, error = _find_chart_yaml(chart_path)
    if error:
        return f"❌ {error}"

    chart_info, error = _parse_chart_yaml(chart_yaml)
    if error:
        return f"❌ {error}"

    chart_dir = chart_info["dir"]
    chart_name = chart_info["name"]
    current_version = chart_info["version"]

    # Calculate new version
    if version:
        new_version = version
    else:
        new_version = _bump_version(current_version, bump)

    # Calculate app version (for Docker tag)
    new_app_version = app_version or new_version

    # Check for Dockerfile
    dockerfile_path = os.path.join(chart_dir, "Dockerfile")
    if not os.path.isfile(dockerfile_path):
        dockerfile_path = os.path.join(os.path.dirname(chart_dir), "Dockerfile")
    has_dockerfile = os.path.isfile(dockerfile_path) and not skip_docker

    # Build summary
    summary = f"""
📦 Helm Deploy Summary
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Chart:        {chart_name}
Path:         {chart_yaml}
Version:      {current_version} → {new_version}
App Version:  {chart_info.get('appVersion', 'N/A')} → {new_app_version}
Docker Build: {'Yes' if has_dockerfile else 'No (no Dockerfile)'}
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Pipeline:
1. Login to ACR
2. {'Build & push Docker image' if has_dockerfile else '(skip Docker - no Dockerfile)'}
3. Update Chart.yaml version
4. Package Helm chart
5. Push to ACR: {ACR_HELM_PREFIX}/{chart_name}
6. Publish to Apollo
"""

    if dry_run:
        return summary + "\n🔍 DRY RUN - No changes made."

    steps.append(summary)

    # Step 2: Login to ACR
    token, error = _get_acr_token()
    if error:
        return f"❌ ACR Login Failed:\n{error}"

    # Docker login
    try:
        docker_login = subprocess.run(
            ["docker", "login", ACR_REGISTRY, "--username", "apollo", "--password-stdin"],
            input=token,
            capture_output=True,
            text=True,
            timeout=30
        )
        if docker_login.returncode == 0:
            steps.append("✓ Docker logged in")
        else:
            steps.append(f"⚠ Docker login failed (continuing): {docker_login.stderr.strip()}")
    except Exception as e:
        steps.append(f"⚠ Docker login skipped: {e}")

    # Helm login
    try:
        helm_login = subprocess.run(
            ["helm", "registry", "login", ACR_REGISTRY, "--username", "apollo", "--password-stdin"],
            input=token,
            capture_output=True,
            text=True,
            timeout=30
        )
        if helm_login.returncode == 0:
            steps.append("✓ Helm logged in")
        else:
            return f"❌ Helm login failed: {helm_login.stderr}"
    except FileNotFoundError:
        return "❌ Helm not found. Please install Helm."

    # Step 3: Build and push Docker image (if applicable)
    if has_dockerfile:
        docker_tag = f"{ACR_DOCKER_PREFIX}/{chart_name}:{new_app_version}"
        steps.append(f"Building Docker image: {docker_tag}")

        try:
            # Build with buildx for multi-arch
            build_result = subprocess.run(
                [
                    "docker", "buildx", "build",
                    "--platform", "linux/amd64,linux/arm64",
                    "--tag", docker_tag,
                    "--output", "type=image,push=true",
                    os.path.dirname(dockerfile_path) or "."
                ],
                capture_output=True,
                text=True,
                timeout=600  # 10 min timeout for builds
            )
            if build_result.returncode == 0:
                steps.append(f"✓ Docker image pushed: {docker_tag}")
            else:
                errors.append(f"Docker build failed: {build_result.stderr}")
                steps.append(f"✗ Docker build failed (continuing with Helm only)")
        except subprocess.TimeoutExpired:
            errors.append("Docker build timed out")
            steps.append("✗ Docker build timed out (continuing with Helm only)")
        except Exception as e:
            errors.append(f"Docker error: {e}")
            steps.append(f"✗ Docker error: {e} (continuing with Helm only)")

    # Step 4: Update Chart.yaml version
    success, error = _update_chart_version(chart_yaml, new_version, new_app_version if has_dockerfile else None)
    if not success:
        return f"❌ Failed to update Chart.yaml: {error}"
    steps.append(f"✓ Updated Chart.yaml to version {new_version}")

    # Step 5: Package Helm chart
    try:
        # Run helm dependency update
        dep_result = subprocess.run(
            ["helm", "dependency", "update"],
            cwd=chart_dir,
            capture_output=True,
            text=True,
            timeout=120
        )
        if dep_result.returncode != 0:
            steps.append(f"⚠ helm dependency update warning: {dep_result.stderr.strip()}")

        # Package
        package_result = subprocess.run(
            ["helm", "package", "."],
            cwd=chart_dir,
            capture_output=True,
            text=True,
            timeout=60
        )
        if package_result.returncode != 0:
            return f"❌ Helm package failed: {package_result.stderr}"

        tgz_file = f"{chart_name}-{new_version}.tgz"
        tgz_path = os.path.join(chart_dir, tgz_file)
        steps.append(f"✓ Packaged: {tgz_file}")
    except Exception as e:
        return f"❌ Helm package error: {e}"

    # Step 6: Push to ACR
    try:
        push_result = subprocess.run(
            ["helm", "push", tgz_path, ACR_HELM_PREFIX],
            capture_output=True,
            text=True,
            timeout=120
        )
        if push_result.returncode != 0:
            return f"❌ Helm push failed: {push_result.stderr}"
        steps.append(f"✓ Pushed to {ACR_HELM_PREFIX}/{chart_name}")

        # Clean up .tgz file
        try:
            os.remove(tgz_path)
        except:
            pass
    except Exception as e:
        return f"❌ Helm push error: {e}"

    # Step 7: Publish to Apollo via apollo-cli
    client_id, client_secret = _get_apollo_credentials()
    oci_url = f"{ACR_HELM_PREFIX}/{chart_name}"
    maven_coordinate = f"{APOLLO_GROUP_ID}:{chart_name}:{new_version}"
    manifest_path = "/tmp/manifest.yml"

    try:
        # Generate manifest.yml using apollo-cli
        init_cmd = [
            "apollo-cli", "product-release", "helm-chart", "init",
            "--maven-coordinate", maven_coordinate,
            "--name", oci_url,
            "--pass-credentials",
            "--repository-url", oci_url,
            "--output-dir", "/tmp/",
            "--username", "apollo",
            "--password", token,
            "--version", new_version,
            "--apollo-client-id", client_id,
            "--apollo-client-secret", client_secret,
            "--apollo-url", APOLLO_URL,
        ]

        init_result = subprocess.run(
            init_cmd,
            capture_output=True,
            text=True,
            timeout=60
        )

        if init_result.returncode != 0:
            errors.append(f"apollo-cli init failed: {init_result.stderr}")
            steps.append(f"⚠ Apollo manifest generation failed: {init_result.stderr.strip()}")
        else:
            steps.append("✓ Generated Apollo manifest")

            # Publish manifest to Apollo
            publish_cmd = [
                "apollo-cli", "publish", "manifest",
                f"--manifest-file={manifest_path}",
                "--apollo-client-id", client_id,
                "--apollo-client-secret", client_secret,
                "--apollo-url", APOLLO_URL,
            ]

            publish_result = subprocess.run(
                publish_cmd,
                capture_output=True,
                text=True,
                timeout=60
            )

            if publish_result.returncode != 0:
                errors.append(f"apollo-cli publish failed: {publish_result.stderr}")
                steps.append(f"⚠ Apollo publish failed: {publish_result.stderr.strip()}")
            else:
                steps.append("✓ Published to Apollo")

    except FileNotFoundError:
        errors.append("apollo-cli not found")
        steps.append("⚠ apollo-cli not found - skipping Apollo publish. Install with: pip install apollo-cli")
    except Exception as e:
        errors.append(f"Apollo error: {e}")
        steps.append(f"⚠ Apollo publish error: {e}")

    # Final summary
    result = "\n".join(steps)
    result += f"\n\n{'🎉' if not errors else '⚠️'} Deployment {'complete' if not errors else 'completed with warnings'}!"
    result += f"\n\nChart: {chart_name}:{new_version}"
    result += f"\nHelm:  {ACR_HELM_PREFIX}/{chart_name}"

    if errors:
        result += f"\n\nWarnings:\n" + "\n".join(f"  - {e}" for e in errors)

    return result


def main():
    mcp.run()


if __name__ == "__main__":
    main()
