"""fix web tool bind address to 0.0.0.0

Revision ID: 2026_05_29_fix_web_tool_bind_address
Revises: 2026_05_29_remove_ssh_keys_mount_from_manifest
Create Date: 2026-05-29 14:00:00.000000

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision: str = "2026_05_29_fix_web_tool_bind_address"
down_revision: Union[str, None] = "2026_05_29_remove_ssh_keys_mount_from_manifest"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _fix_code_server_compose(conn) -> None:
    """Update code-server compose template to bind to 0.0.0.0."""
    result = conn.execute(
        sa.text("""
            SELECT id, compose_template, definition_type
            FROM tool_types
            WHERE name = 'code-server'
        """)
    ).fetchone()

    if result is None:
        return

    tool_id, compose_template, definition_type = result

    if definition_type != "compose" or not compose_template:
        return

    # Fix or add command to bind to 0.0.0.0
    lines = compose_template.split("\n")
    new_lines = []
    image_line_idx = -1
    command_fixed = False
    for i, line in enumerate(lines):
        # Replace broken --bind-addr with correct --host
        if "command:" in line and "--bind-addr" in line:
            indent = line[: len(line) - len(line.lstrip())]
            new_lines.append(f"{indent}command: --host 0.0.0.0")
            command_fixed = True
            continue
        new_lines.append(line)
        if "image:" in line and image_line_idx == -1:
            image_line_idx = i

    # If no command line exists, insert one after image
    if not command_fixed and image_line_idx != -1:
        image_line = lines[image_line_idx]
        indent = image_line[: len(image_line) - len(image_line.lstrip())]
        # Insert after the image line in new_lines
        insert_idx = new_lines.index(image_line) + 1
        new_lines.insert(insert_idx, f"{indent}command: --host 0.0.0.0")
        command_fixed = True

    if not command_fixed:
        return

    updated_compose = "\n".join(new_lines)

    conn.execute(
        sa.text("""
            UPDATE tool_types
            SET compose_template = :compose_template
            WHERE id = :id
        """),
        {"compose_template": updated_compose, "id": tool_id},
    )

    print(f"Updated code-server tool type ({tool_id}) to bind to 0.0.0.0")


def _fix_jupyter_compose(conn) -> None:
    """Update jupyter-notebook compose template to bind to 0.0.0.0."""
    result = conn.execute(
        sa.text("""
            SELECT id, compose_template, definition_type
            FROM tool_types
            WHERE name = 'jupyter-notebook'
        """)
    ).fetchone()

    if result is None:
        return

    tool_id, compose_template, definition_type = result

    if definition_type != "compose" or not compose_template:
        return

    if "command:" in compose_template:
        return

    lines = compose_template.split("\n")
    new_lines = []
    image_line_idx = -1
    for i, line in enumerate(lines):
        new_lines.append(line)
        if "image:" in line and image_line_idx == -1:
            image_line_idx = i
            indent = line[: len(line) - len(line.lstrip())]
            # Jupyter needs --ip=0.0.0.0 to bind to all interfaces
            new_lines.append(
                f"{indent}command: start-notebook.sh --ip=0.0.0.0 --port=8888 --no-browser"
            )

    if image_line_idx == -1:
        return

    updated_compose = "\n".join(new_lines)

    conn.execute(
        sa.text("""
            UPDATE tool_types
            SET compose_template = :compose_template
            WHERE id = :id
        """),
        {"compose_template": updated_compose, "id": tool_id},
    )

    print(f"Updated jupyter-notebook tool type ({tool_id}) to bind to 0.0.0.0:8888")


def upgrade() -> None:
    conn = op.get_bind()
    _fix_code_server_compose(conn)
    _fix_jupyter_compose(conn)


def downgrade() -> None:
    # Cannot safely downgrade without knowing the original compose_template
    pass
