from __future__ import annotations

import base64
import hashlib
import json
import pickle
import platform
import re
from datetime import datetime, timezone
from parser import LogParser
from pathlib import Path
from typing import Generator, Type

from bucket import StructuredLog, bucket
from context import ClinicContext
from dateutil.parser import parse as parsedate
from drain3 import TemplateMiner
from drain3.template_miner_config import TemplateMinerConfig
from task import BaseTask


class LogEntry(StructuredLog):
    eid: str
    ename: str
    level: str
    timestr: str
    file: str
    line: str
    _message: str
    _fields: list[tuple[str, str]]
    component: str
    host: str

    def __init__(
        self,
        xs: tuple[str, str, str, str, str, str, str, list[dict[str, str]]],
        component: str,
        host: str,
    ) -> None:
        self.eid = xs[0]
        self.ename = xs[1]
        self.level = xs[2]
        self.timestr = xs[3]
        self.file = xs[4]
        self.line = xs[5]
        self._message = xs[6]
        self._fields = [(f["Name"], f["Value"]) for f in xs[7]]
        self.component = component
        self.host = host

    def rebuild(self, message: str, fields: list[str]) -> str:
        msg = f"[{self.level}] [{self.file}:{self.line}] [{json.dumps(message)}]"
        fs = ""
        for i, f in enumerate(self.fields()):
            fs += f" [{json.dumps(f[0])}={json.dumps(fields[i])}]"
        msg += fs
        return msg

    def cid(self) -> str:
        fs = "_".join([f[0] for f in self.fields()])
        return f"{self.component}-{self.host}-{self.eid}-{self.file}-{self.line}-{fs}"

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

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


class CombinedUnit(object):
    dt: datetime
    level: str
    eid: str
    ename: str
    component: str
    host: str
    record: str
    count: int

    def __init__(
        self,
        dt: datetime,
        level: str,
        eid: str,
        ename: str,
        component: str,
        host: str,
        record: str,
        count: int,
    ) -> None:
        self.dt = dt.astimezone(timezone.utc)
        self.level = level
        self.eid = eid
        self.ename = ename
        self.component = component
        self.host = host
        self.record = record
        self.count = count


async def shear_log(context: ClinicContext) -> Generator[datetime]:
    with LogParser(Path(__file__).with_name(f"tidb-log-parser-{platform.system().lower()}")) as p:
        docs = []
        async for log in packlogs(context, p):
            timestr = log.dt.strftime("%d/%m/%Y %H:%M:%S.%f")
            hash = hashlib.md5(
                (log.component + log.host + log.eid + timestr + log.record).encode("UTF-8")
            ).hexdigest()
            uid = f"{hash}"
            doc = {
                "_id": uid,
                "eid": log.eid,
                "ename": log.ename or "<unknown>",
                "timestamp": log.dt,
                "level": log.level,
                "host": log.host,
                "component": log.component,
                "message": log.record,
                "count": log.count,
            }
            if len(docs) == 0:
                yield log.dt  # mark the queue head as a checkpoint so that we can resume next time
            docs.append(doc)
            if len(docs) < 1000:
                if len(docs) % 100 == 0:
                    print(timestr, len(docs))
                continue
            await context.write_bulk_logs(docs)
            docs = []
        await context.write_bulk_logs(docs)


async def packlogs(
    context: ClinicContext, parser: LogParser
) -> Generator[CombinedUnit, None, None]:
    logs = []
    async for log in context.get_logs():
        if log[0] not in ["tidb", "pd", "tikv"]:
            continue
        s = parser.parse(log[0], log[2])
        if not s or not s[0]:
            continue
        l = LogEntry(s, log[0], log[1])
        logs.append(l)
        if len(logs) % 1000 == 0:
            print(log[2][:32], len(logs))
        if len(logs) % 10000 == 0:
            for u in process(parsedate(logs[0].timestr), logs):
                yield u
            logs = []
    if len(logs) > 0:
        for u in process(parsedate(logs[0].timestr), logs):
            yield u


def process(dt: datetime, logs: list[LogEntry]) -> Generator[CombinedUnit, None, None]:
    for buc in bucket(logs):
        c = combine(buc)
        l = buc[0]
        record = buc[0].rebuild(c[1], c[2])
        yield CombinedUnit(
            dt=dt,
            level=l.level,
            eid=l.eid,
            ename=l.ename,
            component=l.component,
            host=l.host,
            count=c[0],
            record=record,
        )


def try_max(xs: list[str]) -> str:
    if all([re.match("^[0-9]+$", x) for x in xs]):
        return max([int(x) for x in xs])
    elif all([re.match("^[0-9]+(\.[0-9]+)?$", x) for x in xs]):
        return max([float(x) for x in xs])
    else:
        return max(xs)


def try_min(xs: list[str]) -> str:
    if all([re.match("^[0-9]+$", x) for x in xs]):
        return min([int(x) for x in xs])
    elif all([re.match("^[0-9]+(\.[0-9]+)?$", x) for x in xs]):
        return min([float(x) for x in xs])
    else:
        return min(xs)


def mix(xs: list[list[str]]) -> Generator[tuple[int, list[str]], None, None]:
    tm = TemplateMiner(None, TemplateMinerConfig())
    clusters = {}
    if all([not " " in x[0] for x in xs]):
        minv = try_min([x[0] for x in xs])
        maxv = try_max([x[0] for x in xs])
        if minv != maxv:
            for x in xs:
                x[0] = f"<{minv}-{maxv}>"
    for x in xs:
        r = tm.add_log_message(x[0] if len(x) > 0 else "")
        clusters[r["cluster_id"]] = (
            [x[1:]] if r["cluster_id"] not in clusters else clusters[r["cluster_id"]] + [x[1:]]
        )
    for cid in clusters:
        t = tm.drain.id_to_cluster[cid].get_template()
        if max([len(x) for x in clusters[cid]]) == 0:
            yield (len(clusters[cid]), [t])
        else:
            for m in mix(clusters[cid]):
                yield (m[0], [t] + m[1])


def combine(logs: list[LogEntry]) -> tuple[int, str, list[str]]:
    xs = []
    for l in logs:
        x = [l.message()] + [f[1] for f in l.fields()]
        xs.append(x)
    ls = list(mix(xs))
    assert len(ls) == 1, "params of combine function should belong to same bucket"
    m = ls[0]
    return (m[0], m[1][0], m[1][1:])


class ShearTask(BaseTask):
    index: str
    cluster: str
    package: str
    start_ts: int
    stop_ts: int
    ctx_mark: datetime

    def __init__(
        self,
        index: str,
        cluster: str,
        package: str,
        start_ts: int,
        stop_ts: int,
        **kwargs,
    ) -> None:
        self.index = index
        self.cluster = cluster
        self.package = package
        self.start_ts = start_ts
        self.stop_ts = stop_ts
        self.ctx_mark = None

    @classmethod
    def resume(cls, ctx: str) -> ShearTask:
        task = pickle.loads(base64.b64decode(ctx.encode()))
        task.start_ts = int(task.ctx_mark.timestamp() * 1000)
        return task

    async def run(self) -> None:
        async with ClinicContext(
            "token",
            "user_id",
            self.index,
            self.cluster,
            self.package,
            self.start_ts,
            self.stop_ts,
        ) as ctx:
            async for dt in shear_log(ctx):
                self.ctx_mark = dt

    def suspend(self) -> str:
        if self.ctx_mark == None:
            raise TimeoutError("task timeout before checkpoint")
        return base64.b64encode(pickle.dumps(self)).decode()


def task_class() -> Type[BaseTask]:
    return ShearTask