from __future__ import annotations

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

from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk

# k8s-clinic6j-clinicin-2e7723eccc-70809494.us-west-2.elb.amazonaws.com
ES_HOST = os.environ["ES_HOST"]
API_KEY = os.environ["CLINIC_API_KEY"]


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

    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"]
                if not s["message"].startswith("["):
                    continue
                if self.org_type == "cloud":
                    yield (
                        s["component"],
                        f'{s["host_name"]}-{s["instance"]}',
                        s["message"],
                    )
                else:
                    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}",
                    )

        ts_key = "@timestamp" if self.org_type == "cloud" else "timestamp"
        must = [
            {
                "term": {
                    "user_cluster_id": self.cluster_id,
                }
            },
            {
                "range": {
                    ts_key: {
                        "gte": self.start,
                        "lt": self.stop,
                    }
                },
            },
        ]

        if self.org_type == "cloud":
            self.index = f"filebeat-*"
        else:
            must.append(
                {
                    "term": {
                        "packageid": self.package_id,
                    }
                }
            )

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

        print("searhing logs...")
        resp = await self.source.search(
            index=self.index,
            size=10000,
            query=query,
            sort=ts_key,
            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",
                    cluster_id=self.cluster_id,
                )
                for doc in docs
            ],
        )

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

        self.source = AsyncElasticsearch(
            [
                {
                    "host": ES_HOST,
                    "scheme": "http",
                    "port": 8085,
                    "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()
        await self.sink.close()