from __future__ import annotations

import os
import pickle
import re
import uuid
from datetime import datetime
from typing import Type

import boto3
import numpy as np
import yaml
from context import ClinicContext
from task import BaseTask

S3_BUCKET = os.environ["S3_BUCKET"]  # 'compute-engine-task-resumer'


def condition_match(o, v, line):
    if o == "contains" or o == "contains_many":
        return v in line
    elif o == "regex":
        return bool(re.match(v, line))


def parse_condition(c):
    for o, v in c.items():
        return (o, v)


class Rule:
    def __init__(self, r: dict):
        self._rule = r
        self.matched_conditions = {}

    def process(self, comp, line):
        if self._rule["component"] != comp:
            return
        for i, c in enumerate(self._rule["conditions"]):
            o, v = parse_condition(c)
            if condition_match(o, v, line):
                if i not in self.matched_conditions:
                    self.matched_conditions[i] = 1
                else:
                    self.matched_conditions[i] += 1

    def match(self):
        conditions = []
        for i, c in enumerate(self._rule["conditions"]):
            if i not in self.matched_conditions:
                self.matched_conditions[i] = 0
            o, v = parse_condition(c)
            conditions.append(v)
            if o == "contains_many" and self.matched_conditions[i] < 60:
                return None
            elif self.matched_conditions[i] == 0:
                return None
        return {
            "id": self._rule["id"],
            "name": self._rule["name"],
            "component": self._rule["component"],
            "level": self._rule["level"],
            "type": self._rule["type"],
            "category": self._rule["category"],
            "conditions": conditions,
            "content": {"key": self._rule["advices"]},
            "count": np.min([v for v in self.matched_conditions.values()]).item(),
        }


class LogDiagProcessor:
    def __init__(self):
        with open(os.path.join(os.path.dirname(__file__), "log_diag_rule.yml")) as f:
            rules = yaml.load(f.read(), Loader=yaml.Loader)
        self.rules = [Rule(r) for r in rules]

    def process(self, component, line):
        for r in self.rules:
            r.process(component, line)

    def done(self):
        symptoms = []
        for r in self.rules:
            m = r.match()
            if m is not None:
                symptoms.append(m)
        return symptoms


class DiagTask(BaseTask):
    cluster: str
    start_ts: int
    stop_ts: int
    org_type: str
    ctx_mark: int
    dps: dict
    count: int

    def __init__(self, cluster: str, start_ts: int, stop_ts: int, **kwargs) -> None:
        self.cluster = cluster
        self.start_ts = start_ts * 1000
        self.stop_ts = stop_ts * 1000
        self.org_type = kwargs["org_type"].lower()
        self.org_id = kwargs["org_id"]
        self.packages = kwargs["packages"]
        self.dps = {}
        self.count = 0
        self.ctx_mark = None

    @classmethod
    def resume(cls, ctx: str) -> DiagTask:
        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) -> dict:
        r = []
        if self.org_type != "op" and self.org_type != "cloud":
            raise ValueError("org_type must be op or cloud")
        if isinstance(self.packages, list) == False:
            raise Exception("packages should be a list")
        if len(self.packages) == 0:
            async with ClinicContext(
                self.org_type,
                f"org{self.org_id}_log",
                self.cluster,
                "",
                self.start_ts,
                self.stop_ts,
            ) as ctx:
                async for log in ctx.get_logs():
                    if log == "no log":  # add
                        return {  # add
                            "abnormal": False,  # add
                            "advice": [],  # add
                            "stats": {"scannedLogs": 0},  # add
                        }  # add
                    self.ctx_mark = datetime.strptime(
                        log[2].split("]")[0][1:], "%Y/%m/%d %H:%M:%S.%f +00:00"
                    ).timestamp()
                    inst = log[1]
                    if inst not in self.dps:
                        self.dps[inst] = LogDiagProcessor()
                    self.dps[inst].process(log[0], log[2])
                    self.count += 1
                print(f"{self.count} logs processed")
                for inst, dp in self.dps.items():
                    symptoms = dp.done()
                    r += [{"instance": inst, **sym} for sym in symptoms]
        else:
            for package in self.packages:
                async with ClinicContext(
                    self.org_type,
                    f"org{self.org_id}_log",
                    self.cluster,
                    package,
                    self.start_ts,
                    self.stop_ts,
                ) as ctx:
                    async for log in ctx.get_logs():
                        if log == "no log":
                            return {
                                "abnormal": False,
                                "advice": [],
                                "stats": {"scannedLogs": 0},
                            }
                        self.ctx_mark = datetime.strptime(
                            log[2].split("]")[0][1:],
                            "%Y/%m/%d %H:%M:%S.%f +00:00",
                        ).timestamp()
                        inst = log[1]
                        if inst not in self.dps:
                            self.dps[inst] = LogDiagProcessor()
                        self.dps[inst].process(log[0], log[2])
                        self.count += 1
                    print(f"{self.count} logs processed")
                    for inst, dp in self.dps.items():
                        symptoms = dp.done()
                        r += [{"instance": inst, **sym} for sym in symptoms]
        return {
            "abnormal": len(r) > 0,
            "advice": r,
            "stats": {"scannedLogs": self.count},
        }

    def suspend(self) -> str:
        if self.ctx_mark == None:
            raise TimeoutError("task timeout before checkpoint")
        s3_object = f"logdiag-{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 DiagTask