from __future__ import annotations

import unittest
from abc import ABC, abstractmethod
from typing import Callable

from drain3 import TemplateMiner
from drain3.template_miner_config import TemplateMinerConfig


class StructuredLog(ABC):
    @abstractmethod
    def cid(self) -> str:
        """
        the classification id of the log, should not be unique
        only logs with the same cid may be aggregated together
        """
        pass

    @abstractmethod
    def message(self) -> str:
        pass

    @abstractmethod
    def fields(self) -> list[tuple[str, str]]:
        pass


def split(
    bucket: list[StructuredLog], key: Callable[[StructuredLog], str]
) -> list[list[StructuredLog]]:
    tm = TemplateMiner(None, TemplateMinerConfig())
    singleWorld = list(filter(lambda l: not " " in key(l), bucket))
    bmap = {}
    for log in filter(lambda l: " " in key(l), bucket):
        r = tm.add_log_message(key(log))
        cid = r["cluster_id"]
        if cid not in bmap:
            bmap[cid] = []
        bmap[cid].append(log)
    if len(singleWorld) > 0:
        return [singleWorld] + list(bmap.values())
    return bmap.values()


def mix(
    buckets: list[list[StructuredLog]],
    keys: list[Callable[[StructuredLog], str]],
) -> list[list[StructuredLog]]:
    if len(keys) == 0:
        return buckets
    bs = []
    for bucket in buckets:
        for bkt in split(bucket, keys[0]):
            bs.append(bkt)
    return mix(bs, keys[1:])


def bucket(logs: list[StructuredLog]) -> list[list[StructuredLog]]:
    bmap = {}
    for log in logs:
        cid = log.cid()
        if cid not in bmap:
            bmap[cid] = []
        if len(bmap[cid]) > 0:
            assert len(bmap[cid][0].fields()) == len(
                log.fields()
            ), f"the fields should be same in the same cid({cid})"
        bmap[cid].append(log)
    bs = []
    for b in bmap.values():
        keys = []
        keys.append(lambda x: x.message())
        for i in range(len(b[0].fields())):
            keys.append((lambda idx: lambda x: x.fields()[idx][1])(i))
        bs += mix([b], keys)
    return bs


class MockStructuredLog(StructuredLog):
    _cid: str
    _message: str
    _fields: list[tuple[str, str]]

    def __init__(self, xs: tuple[str, str, list[tuple[str, str]]]):
        self._cid = xs[0]
        self._message = xs[1]
        self._fields = xs[2]

    def __eq__(self, other: MockStructuredLog) -> bool:
        return (
            self._cid == other._cid
            and self._message == other._message
            and self._fields == other._fields
        )

    def cid(self) -> str:
        return self._cid

    def message(self) -> str:
        return self._message

    def fields(self) -> list[tuple[str, str]]:
        return self._fields


class TestSheare(unittest.TestCase):
    def test_bucket(self) -> None:
        cases = [
            (
                [
                    MockStructuredLog(
                        (
                            "cid1",
                            "message 1",
                            [("f1", "foo bar baz"), ("f2", "v is v")],
                        )
                    ),
                    MockStructuredLog(
                        (
                            "cid1",
                            "message 2",
                            [("f1", "aa bb cc"), ("f2", "v is v")],
                        )
                    ),
                    MockStructuredLog(
                        (
                            "cid2",
                            "message 2",
                            [("f1", "v1"), ("f2", "v2"), ("f3", "v3")],
                        )
                    ),
                    MockStructuredLog(
                        (
                            "cid2",
                            "message 3",
                            [("f1", "v3"), ("f2", "v4"), ("f3", "v5")],
                        )
                    ),
                ],
                [
                    [
                        MockStructuredLog(
                            (
                                "cid1",
                                "message 1",
                                [("f1", "foo bar baz"), ("f2", "v is v")],
                            )
                        ),
                    ],
                    [
                        MockStructuredLog(
                            (
                                "cid1",
                                "message 2",
                                [("f1", "aa bb cc"), ("f2", "v is v")],
                            )
                        ),
                    ],
                    [
                        MockStructuredLog(
                            (
                                "cid2",
                                "message 2",
                                [("f1", "v1"), ("f2", "v2"), ("f3", "v3")],
                            )
                        ),
                        MockStructuredLog(
                            (
                                "cid2",
                                "message 3",
                                [("f1", "v3"), ("f2", "v4"), ("f3", "v5")],
                            )
                        ),
                    ],
                ],
            ),
        ]

        for case in cases:
            self.assertEqual(bucket(case[0]), case[1])