from __future__ import annotations

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

import boto3
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 logs import LogEntity, ParsedLog, Parser
from task import BaseTask

S3_BUCKET = os.environ["S3_BUCKET"]


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

    def __init__(self, log: ParsedLog, component: str, host: str) -> None:
        self.eid = log.id
        self.ename = log.name
        self.level = log.level
        self.time = log.time
        self.file = log.file
        self.line = log.line
        self._message = log.message
        self._fields = [(f["Name"], f["Value"]) for f in log.fields]
        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(task: ShearTask, context: ClinicContext) -> Generator[datetime]:
    with Parser() 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 task.stats.get(doc["eid"]) is None:
                task.stats[doc["eid"]] = [
                    doc["ename"],
                    doc["count"],
                    doc["level"],
                    doc["component"],
                ]
                # task.stats[doc["eid"]] = [doc["ename"], doc["count"]]
            else:
                task.stats[doc["eid"]][1] += doc["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)
            # print(docs)
            docs = []
        # await context.write_bulk_logs(docs)
        # print(task.stats)


async def packlogs(context: ClinicContext, parser: Parser) -> Generator[CombinedUnit, None, None]:
    logs = []
    async for log in context.get_logs():
        if log[0] not in ["tidb", "pd", "tikv"]:
            continue
        s = parser.parse(
            LogEntity(
                component=log[0],
                content=log[2],
                time=None,
                cluster=None,
                instance=None,
            )
        )
        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(logs[0].time, logs):
                yield u
            logs = []
    if len(logs) > 0:
        for u in process(logs[0].time, 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
    packages: str
    start_ts: int
    stop_ts: int
    ctx_mark: datetime
    stats: dict

    def __init__(
        self,
        org_id: str,
        cluster: str,
        packages: list,
        start_ts: int,
        stop_ts: int,
        **kwargs,
    ) -> None:
        self.index = f"org{org_id}_log" if kwargs["org_type"] == "op" else f"filebeat-*"
        self.cluster = cluster
        self.packages = packages
        self.start_ts = start_ts * 1000
        self.stop_ts = stop_ts * 1000
        self.ctx_mark = None
        self.stats = {}
        self.org_type = kwargs["org_type"].lower() if kwargs.get("org_type") is not None else None

    @classmethod
    def resume(cls, ctx: str) -> ShearTask:
        s3_object = ctx
        s3 = boto3.client("s3")
        tmpfile = "/tmp/resume.pkl"
        with open(tmpfile, "wb") as f:
            s3.download_fileobj(S3_BUCKET, s3_object, f)
        task = pickle.load(open(tmpfile, "rb"))
        task.start_ts = int(task.ctx_mark * 1000)
        return task

    async def run(self) -> None:
        if self.org_type == "op":
            for package in self.packages:
                async with ClinicContext(
                    "token",
                    "user_id",
                    self.index,
                    self.cluster,
                    package,
                    self.start_ts,
                    self.stop_ts,
                    self.org_type,
                ) as ctx:
                    async for dt in shear_log(self, ctx):
                        self.ctx_mark = dt.timestamp()
                    print(self.stats)
                    err_count = await ctx.get_log_count("ERROR")
                    return {"logStats": self.stats, "errLogCount": err_count}
        if self.org_type == "cloud":
            async with ClinicContext(
                "token",
                "user_id",
                self.index,
                self.cluster,
                "",
                self.start_ts,
                self.stop_ts,
                self.org_type,
            ) as ctx:
                async for dt in shear_log(self, ctx):
                    self.ctx_mark = dt.timestamp()
                print(self.stats)
                err_count = await ctx.get_log_count("ERROR")
                return {"logStats": self.stats, "errLogCount": err_count}

    def suspend(self) -> str:
        if self.ctx_mark == None:
            raise TimeoutError("task timeout before checkpoint")
        s3_object = f"logClustering-{uuid.uuid4()}.pkl"
        s3 = boto3.client("s3")
        with open(f"/tmp/{s3_object}", "wb") as handle:
            pickle.dump(self, handle)
        try:
            s3.upload_file(f"/tmp/{s3_object}", S3_BUCKET, s3_object)
        except Exception as e:
            print(e)
            return False
        return s3_object


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