import json
import uuid

from psycopg2.extensions import cursor as PgCursor
from fixtures.zenith_fixtures import ZenithEnv
from typing import cast

pytest_plugins = ("fixtures.zenith_fixtures")


def helper_compare_branch_list(page_server_cur: PgCursor, env: ZenithEnv, initial_tenant: str):
    """
    Compare branches list returned by CLI and directly via API.
    Filters out branches created by other tests.
    """

    page_server_cur.execute(f'branch_list {initial_tenant}')
    branches_api = sorted(
        map(lambda b: cast(str, b['name']), json.loads(page_server_cur.fetchone()[0])))
    branches_api = [b for b in branches_api if b.startswith('test_cli_') or b in ('empty', 'main')]

    res = env.zenith_cli(["branch"])
    res.check_returncode()
    branches_cli = sorted(map(lambda b: b.split(':')[-1].strip(), res.stdout.strip().split("\n")))
    branches_cli = [b for b in branches_cli if b.startswith('test_cli_') or b in ('empty', 'main')]

    res = env.zenith_cli(["branch", f"--tenantid={initial_tenant}"])
    res.check_returncode()
    branches_cli_with_tenant_arg = sorted(
        map(lambda b: b.split(':')[-1].strip(), res.stdout.strip().split("\n")))
    branches_cli_with_tenant_arg = [
        b for b in branches_cli if b.startswith('test_cli_') or b in ('empty', 'main')
    ]

    assert branches_api == branches_cli == branches_cli_with_tenant_arg


def test_cli_branch_list(zenith_simple_env: ZenithEnv):
    env = zenith_simple_env
    page_server_conn = env.pageserver.connect()
    page_server_cur = page_server_conn.cursor()

    # Initial sanity check
    helper_compare_branch_list(page_server_cur, env, env.initial_tenant)

    # Create a branch for us
    res = env.zenith_cli(["branch", "test_cli_branch_list_main", "empty"])
    assert res.stderr == ''
    helper_compare_branch_list(page_server_cur, env, env.initial_tenant)

    # Create a nested branch
    res = env.zenith_cli(["branch", "test_cli_branch_list_nested", "test_cli_branch_list_main"])
    assert res.stderr == ''
    helper_compare_branch_list(page_server_cur, env, env.initial_tenant)

    # Check that all new branches are visible via CLI
    res = env.zenith_cli(["branch"])
    assert res.stderr == ''
    branches_cli = sorted(map(lambda b: b.split(':')[-1].strip(), res.stdout.strip().split("\n")))

    assert 'test_cli_branch_list_main' in branches_cli
    assert 'test_cli_branch_list_nested' in branches_cli


def helper_compare_tenant_list(page_server_cur: PgCursor, env: ZenithEnv):
    page_server_cur.execute(f'tenant_list')
    tenants_api = sorted(
        map(lambda t: cast(str, t['id']), json.loads(page_server_cur.fetchone()[0])))

    res = env.zenith_cli(["tenant", "list"])
    assert res.stderr == ''
    tenants_cli = sorted(map(lambda t: t.split()[0], res.stdout.splitlines()))

    assert tenants_api == tenants_cli


def test_cli_tenant_list(zenith_simple_env: ZenithEnv):
    env = zenith_simple_env
    page_server_conn = env.pageserver.connect()
    page_server_cur = page_server_conn.cursor()

    # Initial sanity check
    helper_compare_tenant_list(page_server_cur, env)

    # Create new tenant
    tenant1 = uuid.uuid4().hex
    res = env.zenith_cli(["tenant", "create", tenant1])
    res.check_returncode()

    # check tenant1 appeared
    helper_compare_tenant_list(page_server_cur, env)

    # Create new tenant
    tenant2 = uuid.uuid4().hex
    res = env.zenith_cli(["tenant", "create", tenant2])
    res.check_returncode()

    # check tenant2 appeared
    helper_compare_tenant_list(page_server_cur, env)

    res = env.zenith_cli(["tenant", "list"])
    res.check_returncode()
    tenants = sorted(map(lambda t: t.split()[0], res.stdout.splitlines()))

    assert env.initial_tenant in tenants
    assert tenant1 in tenants
    assert tenant2 in tenants