import unittest
from io import StringIO
from unittest.mock import patch

from o11ydiagnosis.subtasks.slowlogconverter.slowlog import StreamParser
from o11ydiagnosis.subtasks.slowlogconverter.write_deltalake import (
    slowlogs_schema,
    write_slowlogs_to_deltalake,
)

READ_POOL_TASK_DETAILS = "{tasks:1, poll_count:{total:2, avg:2, max:2, min:2}, dispatch_count:{total:2, max:2, min:2}, task_wall_time:{total:611.5µs, avg:611.5µs, max:611.5µs, min:611.5µs}, queue_wait:{total:25.8µs, avg:12.9µs, max:13.9µs, min:12µs}}"

TEST_LOG_1 = (
    """2022-08-31T07:12:54.233489375Z;ip-10-250-24-4.us-west-2.compute.internal;tidb61992;db-tidb-0;slowlog # Time: 2022-08-31T07:12:54.096687178Z
# Txn_start_ts: 0
# User@Host: root[root] @ 10.250.24.175 [10.250.24.175]
# Conn_ID: 6204412174144681481
# Query_time: 1.000883579
# Parse_time: 0.000046967
# Compile_time: 0.000166604
# Rewrite_time: 0.000063231
# Optimize_time: 0.000029576
# Wait_TS: 0
# Read_pool_task_details: """
    + READ_POOL_TASK_DETAILS
    + """
# Is_internal: false
# Digest: 926597c5c848f662ee6caba21f181a5aac72d8b2aee8ebbcc8ee361b46f7d9cb
# Num_cop_tasks: 0
# Prepared: false
# Plan_from_cache: false
# Plan_from_binding: false
# Has_more_results: false
# KV_total: 0
# PD_total: 0.000006361
# Backoff_total: 0
# Write_sql_response_total: 0.000000887
# Result_rows: 1
# Succ: true
# IsExplicitTxn: false
# Plan: tidb_decode_plan('ffBhMAkzXzMJMAkxCXNsZWVwKDEpLT5Db2x1bW4jMQkxCXRpbWU6MXMsIGxvb3BzOjIsIENvbmN1cnJlbmN5Ok9GRgkwIEJ5dGVzCU4vQQoxCTI1XzQJMAkxCXJvd3M6MQkxCXQBQgwyOThuGUUgCU4vQQlOL0EK')
# Plan_digest: a55eaee1d30304ad05500afbbfff1409ce22376e0f062c0d06ff34b8ed9c2b47
select sleep (?);"""
)


class Test(unittest.TestCase):
    def test_parse(self):
        parser = StreamParser(reader=StringIO(TEST_LOG_1))
        sl = parser.read_next()
        self.assertIsNotNone(sl)
        self.assertEqual(sl.read_pool_task_details, READ_POOL_TASK_DETAILS)
        self.assertEqual(sl.to_dict()["Read_pool_task_details"], READ_POOL_TASK_DETAILS)
        sl = parser.read_next()
        self.assertIsNone(sl)

    @patch("o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.try_to_optimize_table")
    @patch("o11ydiagnosis.subtasks.slowlogconverter.write_deltalake.write_to_table")
    def test_write_read_pool_task_details_to_delta_schema(self, write_to_table, _):
        parser = StreamParser(reader=StringIO(TEST_LOG_1))
        slowlog = parser.read_next()
        slowlog.id = 1

        write_slowlogs_to_deltalake(
            [slowlog],
            region="us-west-2",
            cluster="test-cluster",
            base_path="s3://test-bucket/deltalake/test-org/test-cluster",
            optimize=False,
            partition_date="20260821",
        )

        field = slowlogs_schema.field("read_pool_task_details")
        self.assertTrue(field.nullable)
        self.assertEqual(str(field.type), "string")

        table = write_to_table.call_args.kwargs["source_pyarrow_table"]
        self.assertEqual(
            table.column("read_pool_task_details").to_pylist(),
            [READ_POOL_TASK_DETAILS],
        )


if __name__ == "__main__":
    unittest.main()