from __future__ import annotations

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

from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk

# clinic-staging-133559976.cn-north-1.elb.amazonaws.com.cn
ES_HOST = os.environ["ES_HOST"]
API_KEY = os.environ["CLINIC_API_KEY"]


class ClinicContext:
    def __init__(self, clinic_token, user_id, index, cluster_id, package_id, start, stop):
        self.index = index
        self.cluster_id = cluster_id
        self.package_id = package_id
        self.clinic_token = clinic_token
        self.user_id = user_id
        self.start = start
        self.stop = stop

    def whoami(self):
        return self.user_id

    def get_metrics(self):
        return []

    async def get_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"]
                d = datetime.fromtimestamp(s["timestamp"] / 1000, timezone.utc)
                level = s["level"].upper()
                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}",
                )

        query = {
            "bool": {
                "must": {"match_all": {}},
                "filter": {
                    "bool": {
                        "must_not": [{"term": {"component": "tiflash"}}],
                        "must": [
                            {
                                "term": {
                                    "cluster_id": self.cluster_id,
                                }
                            },
                            {
                                "term": {
                                    "packageid": self.package_id,
                                }
                            },
                            {
                                "range": {
                                    "timestamp": {
                                        "gte": self.start,
                                        "lt": self.stop,
                                    }
                                },
                            },
                        ],
                    }
                },
            }
        }
        print("searhing logs...")
        resp = await self.source.search(
            index=self.index,
            size=10000,
            query=query,
            sort="timestamp",
            scroll="60s",
        )
        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 write_bulk_logs(self, docs):
        await async_bulk(
            self.sink,
            [
                dict(
                    doc,
                    _index="clinic-log-analysis",
                    source=self.index,
                    cluster_id=self.cluster_id,
                )
                for doc in docs
            ],
        )

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

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