import json
from datetime import date, datetime
from unittest.mock import Mock, patch

import pytest

from .client import (
    _ATHENA_API_SLEEP_SECONDS,
    GlueAthenaClient,
    _column,
    _format_table_description,
    _iter_table_columns,
    _parse_table,
    _query_from_execution,
    _query_user_from_cloudtrail_event,
    _s3_to_http,
    _table_type,
    _warehouse_user,
)


@pytest.mark.parametrize(
    ("url", "expected"),
    [
        ("https://example.com/path", "https://example.com/path"),
        (
            "s3://my-bucket/path/to/object.csv",
            "https://s3.console.aws.amazon.com/s3/buckets/"
            "my-bucket/?prefix=path/to/object.csv",
        ),
        (
            "s3://my-bucket",
            "https://s3.console.aws.amazon.com/s3/buckets/my-bucket/?prefix=",
        ),
        (
            "s3://my-bucket/folder/",
            "https://s3.console.aws.amazon.com/s3/buckets/"
            "my-bucket/?prefix=folder/",
        ),
    ],
)
def test__s3_to_http(url: str, expected: str):
    assert _s3_to_http(url) == expected


def test__format_table_description_without_description():
    location = "s3://bucket/folder/"
    expected = (
        "Bucket location: [s3://bucket/folder/]"
        "(https://s3.console.aws.amazon.com/s3/buckets/"
        "bucket/?prefix=folder/)"
    )
    assert _format_table_description(location, None) == expected


def test__format_table_description_with_description():
    location = "s3://bucket/folder/"
    description = "Table description"
    expected = (
        "Table description\n"
        "Bucket location: [s3://bucket/folder/]"
        "(https://s3.console.aws.amazon.com/s3/buckets/"
        "bucket/?prefix=folder/)"
    )
    assert _format_table_description(location, description) == expected


def test__format_table_description_empty_location_without_description():
    assert _format_table_description("", None) is None


def test__format_table_description_empty_location_with_description():
    assert (
        _format_table_description("", "Table description")
        == "Table description"
    )


def test__table_type():
    assert _table_type(None) == "TABLE"
    assert _table_type("EXTERNAL_TABLE") == "EXTERNAL"
    assert _table_type("VIRTUAL_VIEW") == "VIEW"
    assert _table_type("random") == "TABLE"


def test__parse_column_with_name():
    parsed = _column(
        "table_arn",
        "col_a",
        "varchar",
        "desc",
        2,
    )
    assert parsed == {
        "column_id": "table_arn.col_a",
        "column_name": "col_a",
        "data_type": "varchar",
        "description": "desc",
        "ordinal_position": 2,
        "table_id": "table_arn",
        "tags": [],
    }


def test__resource_arn():
    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._region = "eu-west-1"
    client._account_id = "123"

    assert client._resource_arn("database", "db1") == (
        "arn:aws:glue:eu-west-1:123:database/db1"
    )
    assert client._resource_arn("table", "db1/table1") == (
        "arn:aws:glue:eu-west-1:123:table/db1/table1"
    )


def test_databases_is_hard_coded_catalog():
    client = GlueAthenaClient.__new__(GlueAthenaClient)
    assert client.databases() == [
        {"database_id": "awsdatacatalog", "database_name": "awsdatacatalog"}
    ]


def test_view_ddl_extracts_only_views():
    paginator = Mock()
    paginator.paginate.side_effect = [
        [
            {
                "TableList": [
                    {
                        "Name": "orders_view",
                        "TableType": "VIRTUAL_VIEW",
                        "ViewOriginalText": "SELECT * FROM orders",
                    },
                    {
                        "Name": "orders",
                        "TableType": "EXTERNAL_TABLE",
                    },
                ]
            }
        ],
        [
            {
                "TableList": [
                    {
                        "Name": "customers_view",
                        "TableType": "VIRTUAL_VIEW",
                        "ViewDefinition": {
                            "Representations": [
                                {
                                    "Dialect": "ATHENA",
                                    "ViewOriginalText": "SELECT * FROM customers",
                                }
                            ]
                        },
                    }
                ]
            }
        ],
    ]
    glue_client = Mock()
    glue_client.get_paginator.return_value = paginator

    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._glue_client = glue_client
    client.schemas = Mock(
        return_value=[
            {"schema_name": "analytics"},
            {"schema_name": "sales"},
        ]
    )

    assert client.view_ddl() == [
        {
            "database_name": "awsdatacatalog",
            "schema_name": "analytics",
            "view_definition": (
                "CREATE VIEW analytics.orders_view AS\nSELECT * FROM orders"
            ),
            "view_name": "orders_view",
        },
        {
            "database_name": "awsdatacatalog",
            "schema_name": "sales",
            "view_definition": (
                "CREATE VIEW sales.customers_view AS\nSELECT * FROM customers"
            ),
            "view_name": "customers_view",
        },
    ]


def test_users_deduplicates_query_log_identities():
    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._query_users = Mock(
        return_value={
            "q-1": {
                "user_email": None,
                "user_id": "user-2",
                "user_name": None,
            },
            "q-2": {
                "user_email": None,
                "user_id": "user-1",
                "user_name": "alice",
            },
            "q-3": {
                "user_email": None,
                "user_id": "user-2",
                "user_name": "bob",
            },
            "q-4": {
                "user_email": None,
                "user_id": None,
                "user_name": "ignored",
            },
        }
    )

    users = client.users(date(2026, 3, 18))

    client._query_users.assert_called_once_with(
        datetime(2026, 3, 12),
        datetime(2026, 3, 19),
    )

    assert users == [
        {
            "user_email": None,
            "user_id": "user-2",
            "user_name": "bob",
        },
        {
            "user_email": None,
            "user_id": "user-1",
            "user_name": "alice",
        },
    ]


def test__warehouse_user_skips_missing_identifier():
    assert (
        _warehouse_user(
            {
                "user_email": None,
                "user_id": None,
                "user_name": "alice",
            }
        )
        is None
    )


def test__warehouse_user_falls_back_to_user_id_for_name():
    assert _warehouse_user(
        {
            "user_email": None,
            "user_id": "ABCD:session",
            "user_name": None,
        }
    ) == {
        "user_email": None,
        "user_id": "ABCD:session",
        "user_name": "ABCD:session",
    }


def test__query_user_from_cloudtrail_event_reads_query_id_and_user():
    event = {
        "CloudTrailEvent": json.dumps(
            {
                "responseElements": {"queryExecutionId": "q-1"},
                "userIdentity": {
                    "principalId": "AROATESTPRINCIPAL:alice.analyst@example.com",
                    "userName": "team-role",
                },
            }
        )
    }

    assert _query_user_from_cloudtrail_event(event) == (
        "q-1",
        {
            "user_email": "alice.analyst@example.com",
            "user_id": "alice.analyst@example.com",
            "user_name": "alice.analyst@example.com",
        },
    )


def test__query_user_from_cloudtrail_event_skips_missing_query_id():
    event = {
        "CloudTrailEvent": json.dumps(
            {
                "responseElements": None,
                "userIdentity": {
                    "principalId": "ABCD:session",
                    "userName": "alice",
                },
            }
        )
    }

    assert _query_user_from_cloudtrail_event(event) is None


def test__query_from_execution_uses_query_user():
    query_execution = {
        "Query": "SELECT 1",
        "QueryExecutionContext": {"Database": "analytics"},
        "QueryExecutionId": "q-1",
        "Status": {
            "CompletionDateTime": datetime(2026, 3, 18, 10, 5),
            "SubmissionDateTime": datetime(2026, 3, 18, 10, 0),
        },
    }

    assert _query_from_execution(
        query_execution,
        execution_date=date(2026, 3, 18),
        query_user={
            "user_email": "alice.analyst@example.com",
            "user_id": "alice.analyst@example.com",
            "user_name": "alice.analyst@example.com",
        },
    ) == {
        "database_name": "analytics",
        "end_time": "2026-03-18T10:05:00",
        "query_id": "q-1",
        "query_text": "SELECT 1",
        "start_time": "2026-03-18T10:00:00",
        "user_email": "alice.analyst@example.com",
        "user_id": "alice.analyst@example.com",
        "user_name": "alice.analyst@example.com",
    }


@patch(
    "source.packages.extractor.castor_extractor.warehouse.glue_athena.client.time.sleep"
)
def test__iter_cloudtrail_query_events_paginates(sleep_mock):
    cloudtrail_client = Mock()
    cloudtrail_client.lookup_events.side_effect = [
        {
            "Events": [{"CloudTrailEvent": "{}"}],
            "NextToken": "next",
        },
        {
            "Events": [{"CloudTrailEvent": "{}"}],
        },
    ]

    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._cloudtrail_client = cloudtrail_client

    events = client._iter_cloudtrail_query_events(
        datetime(2026, 3, 18),
        datetime(2026, 3, 19),
    )

    assert events == [
        {"CloudTrailEvent": "{}"},
        {"CloudTrailEvent": "{}"},
    ]
    sleep_mock.assert_called_once_with(_ATHENA_API_SLEEP_SECONDS)


def test_workgroups_filters_enabled_and_paginates():
    athena_client = Mock()
    athena_client.can_paginate.return_value = True
    paginator = Mock()
    paginator.paginate.return_value = [
        {
            "WorkGroups": [
                {"Name": "primary", "State": "ENABLED"},
                {"Name": "disabled", "State": "DISABLED"},
            ],
        },
        {
            "WorkGroups": [{"Name": "analytics", "State": "ENABLED"}],
        },
    ]
    athena_client.get_paginator.return_value = paginator

    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._athena_client = athena_client

    assert client.workgroups() == ["primary", "analytics"]


@patch(
    "source.packages.extractor.castor_extractor.warehouse.glue_athena.client.time.sleep"
)
def test_workgroups_falls_back_to_manual_pagination(sleep_mock):
    athena_client = Mock()
    athena_client.can_paginate.return_value = False
    athena_client.list_work_groups.side_effect = [
        {
            "WorkGroups": [
                {"Name": "primary", "State": "ENABLED"},
                {"Name": "disabled", "State": "DISABLED"},
            ],
            "NextToken": "next",
        },
        {
            "WorkGroups": [{"Name": "analytics", "State": "ENABLED"}],
        },
    ]

    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._athena_client = athena_client

    assert client.workgroups() == ["primary", "analytics"]
    sleep_mock.assert_called_once_with(_ATHENA_API_SLEEP_SECONDS)


@patch(
    "source.packages.extractor.castor_extractor.warehouse.glue_athena.client.time.sleep"
)
def test__query_users_reads_cloudtrail_events(sleep_mock):
    cloudtrail_client = Mock()
    cloudtrail_client.lookup_events.side_effect = [
        {
            "Events": [
                {
                    "CloudTrailEvent": json.dumps(
                        {
                            "responseElements": {"queryExecutionId": "q-1"},
                            "userIdentity": {
                                "principalId": "ABCD:session",
                                "userName": "alice",
                            },
                        }
                    )
                }
            ],
            "NextToken": "next",
        },
        {
            "Events": [
                {
                    "CloudTrailEvent": json.dumps(
                        {
                            "responseElements": {"queryExecutionId": "q-2"},
                            "userIdentity": {
                                "arn": "arn:aws:sts::123:assumed-role/demo/session",
                            },
                        }
                    )
                },
                {
                    "CloudTrailEvent": json.dumps(
                        {
                            "responseElements": {"queryExecutionId": "q-3"},
                            "userIdentity": {
                                "principalId": (
                                    "AROATESTPRINCIPAL:"
                                    "alice.analyst@example.com"
                                ),
                                "userName": "team-role",
                            },
                        }
                    )
                },
            ]
        },
    ]

    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._cloudtrail_client = cloudtrail_client

    users = client._query_users(
        datetime(2026, 3, 18),
        datetime(2026, 3, 19),
    )

    assert users == {
        "q-1": {
            "user_email": None,
            "user_id": "ABCD:session",
            "user_name": "ABCD:session",
        },
        "q-2": {
            "user_email": None,
            "user_id": "arn:aws:sts::123:assumed-role/demo/session",
            "user_name": "arn:aws:sts::123:assumed-role/demo/session",
        },
        "q-3": {
            "user_email": "alice.analyst@example.com",
            "user_id": "alice.analyst@example.com",
            "user_name": "alice.analyst@example.com",
        },
    }
    sleep_mock.assert_called_once_with(_ATHENA_API_SLEEP_SECONDS)


def test__query_users_skips_null_response_elements():
    cloudtrail_client = Mock()
    cloudtrail_client.lookup_events.return_value = {
        "Events": [
            {
                "CloudTrailEvent": json.dumps(
                    {
                        "responseElements": None,
                        "userIdentity": {
                            "principalId": "ABCD:session",
                            "userName": "alice",
                        },
                    }
                )
            }
        ]
    }

    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._cloudtrail_client = cloudtrail_client

    assert (
        client._query_users(
            datetime(2026, 3, 18),
            datetime(2026, 3, 19),
        )
        == {}
    )


@patch(
    "source.packages.extractor.castor_extractor.warehouse.glue_athena.client.time.sleep"
)
def test__workgroup_queries_filters_target_day_and_stops_early(sleep_mock):
    athena_client = Mock()
    athena_client.list_query_executions.return_value = {
        "QueryExecutionIds": ["q-new", "q-target", "q-old"],
    }
    athena_client.batch_get_query_execution.return_value = {
        "QueryExecutions": [
            {
                "Query": "SELECT * FROM too_new",
                "QueryExecutionId": "q-new",
                "Status": {
                    "CompletionDateTime": datetime(2026, 3, 19, 9, 5),
                    "SubmissionDateTime": datetime(2026, 3, 19, 9, 0),
                },
            },
            {
                "Query": "SELECT * FROM target",
                "QueryExecutionContext": {"Database": "analytics"},
                "QueryExecutionId": "q-target",
                "Status": {
                    "CompletionDateTime": datetime(2026, 3, 18, 10, 5),
                    "SubmissionDateTime": datetime(2026, 3, 18, 10, 0),
                },
            },
            {
                "Query": "SELECT * FROM too_old",
                "QueryExecutionId": "q-old",
                "Status": {
                    "CompletionDateTime": datetime(2026, 3, 17, 8, 5),
                    "SubmissionDateTime": datetime(2026, 3, 17, 8, 0),
                },
            },
        ]
    }

    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._athena_client = athena_client

    queries = client._workgroup_queries(
        "primary",
        date(2026, 3, 18),
        users_by_query_id={
            "q-target": {
                "user_email": None,
                "user_id": "user-1",
                "user_name": "alice",
            }
        },
    )

    assert queries == [
        {
            "database_name": "analytics",
            "end_time": "2026-03-18T10:05:00",
            "query_id": "q-target",
            "query_text": "SELECT * FROM target",
            "start_time": "2026-03-18T10:00:00",
            "user_email": None,
            "user_id": "user-1",
            "user_name": "alice",
        },
    ]
    sleep_mock.assert_not_called()


def test_queries_uses_provided_workgroups():
    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client._query_users = Mock(return_value={"q-1": {"user_id": "user-1"}})
    client._workgroup_queries = Mock(
        side_effect=[[{"query_id": "q-1"}], [{"query_id": "q-2"}]]
    )

    queries = client.queries(
        date(2026, 3, 18),
        workgroups=["primary", "analytics"],
    )

    assert queries == [{"query_id": "q-1"}, {"query_id": "q-2"}]


def test_queries_keeps_results_when_cloudtrail_enrichment_fails():
    client = GlueAthenaClient.__new__(GlueAthenaClient)
    client.workgroups = Mock(return_value=["primary"])
    client._query_users = Mock(side_effect=RuntimeError("boom"))
    client._workgroup_queries = Mock(return_value=[{"query_id": "q-1"}])

    with patch(
        "source.packages.extractor.castor_extractor.warehouse.glue_athena.client.logger"
    ) as logger_mock:
        queries = client.queries(date(2026, 3, 18))

    assert queries == [{"query_id": "q-1"}]
    logger_mock.warning.assert_called_once()


def test__parse_table():
    raw_table = {
        "Name": "orders",
        "Description": "Orders table",
        "TableType": "EXTERNAL_TABLE",
        "StorageDescriptor": {"Location": "s3://bucket/orders/"},
    }

    parsed = _parse_table(raw_table, "schema-id", "table-arn")
    assert parsed == {
        "table_id": "table-arn",
        "schema_id": "schema-id",
        "table_name": "orders",
        "description": (
            "Orders table\n"
            "Bucket location: [s3://bucket/orders/]"
            "(https://s3.console.aws.amazon.com/s3/buckets/"
            "bucket/?prefix=orders/)"
        ),
        "tags": [],
        "type": "EXTERNAL",
    }


def test__iter_table_columns_keeps_order_with_flattening():
    raw_table = {
        "StorageDescriptor": {
            "Columns": [
                {"Name": "id", "Type": "int"},
                {"Name": "ds", "Type": "string"},
            ]
        },
        "PartitionKeys": [
            {"Name": "ds", "Type": "string"},
            {"Name": "country", "Type": "varchar"},
        ],
    }

    parsed = _iter_table_columns(raw_table, "table-arn")
    assert [column["column_name"] for column in parsed] == [
        "id",
        "ds",
        "country",
    ]
    assert [column["ordinal_position"] for column in parsed] == [0, 1, 2]


def test__iter_table_columns_flattens_nested_types():
    raw_table = {
        "Parameters": {},
        "StorageDescriptor": {
            "Columns": [
                {
                    "Name": "activity",
                    "Type": "struct<zone:struct<city:string>,radius:int>",
                }
            ]
        },
        "PartitionKeys": [],
    }

    parsed = _iter_table_columns(raw_table, "table-arn")
    assert [column["column_name"] for column in parsed] == [
        "activity",
        "activity.zone",
        "activity.zone.city",
        "activity.radius",
    ]
    assert [column["data_type"] for column in parsed] == [
        "struct",
        "struct",
        "string",
        "int",
    ]


def test__iter_table_columns_sets_spark_descriptions():
    raw_table = {
        "Parameters": {
            "spark.sql.sources.schema.part.0": (
                '{"type":"struct","fields":['
                '{"name":"activity","type":{"type":"struct","fields":['
                '{"name":"zone","type":{"type":"struct","fields":['
                '{"name":"city","type":"string","nullable":true,'
                '"metadata":{"comment":"city comment"}}'
                ']},"nullable":true,"metadata":{"comment":"zone comment"}},'
                '{"name":"radius","type":"integer","nullable":true,'
                '"metadata":{"comment":"radius comment"}}'
                ']},"nullable":true,"metadata":{"comment":"activity comment"}}'
                "]}"
            )
        },
        "StorageDescriptor": {
            "Columns": [
                {
                    "Name": "activity",
                    "Type": "struct<zone:struct<city:string>,radius:int>",
                }
            ]
        },
        "PartitionKeys": [],
    }

    parsed = _iter_table_columns(raw_table, "table-arn")
    assert [column["description"] for column in parsed] == [
        "activity comment",
        "zone comment",
        "city comment",
        "radius comment",
    ]


def test__iter_table_columns_prefers_spark_description_over_glue_comment():
    raw_table = {
        "Parameters": {
            "spark.sql.sources.schema.part.0": (
                '{"type":"struct","fields":['
                '{"name":"activity","type":"string","nullable":true,'
                '"metadata":{"comment":"spark comment"}}'
                "]}"
            )
        },
        "StorageDescriptor": {
            "Columns": [
                {
                    "Comment": "glue comment",
                    "Name": "activity",
                    "Type": "string",
                }
            ]
        },
        "PartitionKeys": [],
    }

    parsed = _iter_table_columns(raw_table, "table-arn")
    assert [column["description"] for column in parsed] == ["spark comment"]
