from __future__ import annotations

import asyncio
import os
from datetime import datetime, timezone
from typing import Generator

from elasticsearch import AsyncElasticsearch, exceptions
from logs import Client, Parser

ES_HOST = os.environ["ES_HOST"]
API_KEY = os.environ["CLINIC_API_KEY"]
ES_PORT = os.environ["ES_PORT"]


class ClinicContext:
    def __init__(
        self,
        token,
        user_id,
        index,
        cluster_id,
        package_id,
        start,
        stop,
        org_type,
    ):
        self.org_type = org_type
        self.index = index
        self.cluster_id = cluster_id
        self.package_id = package_id
        self.start = start
        self.stop = stop
        self.level_stats = {}

    async def get_logs(self) -> Generator[tuple[str, str, str], None, None]:
        if self.org_type == "cloud":
            g = self.get_cloud_logs()
        else:
            g = self.get_op_logs()
        async for x in g:
            yield x

    async def get_cloud_logs(
        self,
    ) -> Generator[tuple[str, str, str], None, None]:
        client = Client()
        parser = Parser()
        for entity in client.pull(
            int(self.start / 1000),
            int(self.stop / 1000),
            cluster_ids=[int(self.cluster_id)],
            components=["tidb", "tikv", "pd"],
        ):
            level = parser.parse(entity).level.upper()
            if level not in self.level_stats:
                self.level_stats[level] = 0
            self.level_stats[level] += 1
            await asyncio.sleep(0)  # make sure the job can be suspended
            yield (entity.component, entity.instance, entity.content)

    async def get_op_logs(self) -> Generator[tuple[str, str, str], None, None]:
        def handle_data(
            hits: list[dict],
        ) -> Generator[tuple[str, str, str], None, None]:
            for hit in hits:
                s = hit["_source"]
                if not s["message"].startswith("["):
                    continue
                d = datetime.fromtimestamp(s["timestamp"] / 1000, timezone.utc)
                level = s["level"].upper()
                if level not in self.level_stats:
                    self.level_stats[level] = 0
                self.level_stats[level] += 1
                message = s["message"]
                timestr = f"{d.year}/{d.month:02d}/{d.day:02d} {d.hour:02d}:{d.minute:02d}:{d.second:02d}.{int(d.microsecond/1000):03d} +00:00"
                yield (
                    s["component"],
                    f'{s["host"]}-{s["instanceid"]}',
                    f"[{timestr}] [{level}] {message}",
                )

        ts_key = "timestamp"
        query = self.generate_query(ts_key)
        print("searhing logs...")
        try:
            resp = await self.source.search(
                index=self.index,
                size=10000,
                query=query,
                sort=ts_key,
                scroll="60s",
            )
        except exceptions.NotFoundError:
            for l in ["no log"]:
                yield l
        s = resp["_shards"]
        print(
            f'total: {s["total"]}, successful: {s["successful"]}, skipped: {s["skipped"]}, failed: {s["failed"]}, hits: {len(resp["hits"]["hits"])}'
        )
        while len(resp["hits"]["hits"]) > 0:
            for l in handle_data(resp["hits"]["hits"]):
                yield l
            print("scrolling logs...")
            resp = await self.source.scroll(scroll_id=resp["_scroll_id"], scroll="60s")
            s = resp["_shards"]
            print(
                f'total: {s["total"]}, successful: {s["successful"]}, skipped: {s["skipped"]}, failed: {s["failed"]}, hits: {len(resp["hits"]["hits"])}'
            )

    async def __aenter__(self) -> ClinicContext:
        self.source = AsyncElasticsearch(
            [
                {
                    "host": ES_HOST,
                    "scheme": "http",
                    "port": int(ES_PORT),
                    "path_prefix": "/clinic/api/v1/data/log",
                }
            ],
            request_timeout=600,
            headers={
                "Authorization": f"Bearer {API_KEY}",
                "X-OrgType": self.org_type,
            },
        )
        return self

    async def __aexit__(self, *args) -> None:
        await self.source.close()

    def generate_query(self, ts_key: str, filter=None) -> dict:
        must = [
            {
                "term": {
                    "cluster_id": self.cluster_id,
                }
            },
            {
                "range": {
                    ts_key: {
                        "gte": self.start,
                        "lt": self.stop,
                    }
                },
            },
            {
                "term": {
                    "packageid": self.package_id,
                }
            },
        ]

        if isinstance(filter, dict):
            must.append({"term": filter})

        query = {
            "bool": {
                "must": {"match_all": {}},
                "filter": {
                    "bool": {
                        "must_not": [{"term": {"component": "tiflash"}}],
                        "must": must,
                    }
                },
            }
        }
        return query

    async def get_log_count(self, level) -> int:
        return 0 if level not in self.level_stats else self.level_stats[level]