import os
import sys
import time
import logging
from datetime import datetime
from o11ydiagnosis.subtasks.slowlogconverter.task import read_keys_from_merged_logs
from o11ydiagnosis.subtasks.slowlogconverter.objstore import S3, GCS, AzureBlob

_h = logging.StreamHandler(sys.stdout)
_h.setFormatter(logging.Formatter("%(asctime)s [%(name)s] [%(levelname)s] %(message)s"))
log = logging.Logger("regionscompactor", logging.DEBUG)
log.addHandler(_h)


class Config:
    cluster: str
    interval_secs: int

    def __init__(self, task):
        self.env = task.env
        self.cluster_id = task.cluster_id
        self.o11y_bucket = o11y_s3_bucket(task.env, task.region, task.serverless)
        self.o11y_prefix = o11y_s3_prefix(task.tenant_id, task.project_id, task.cluster_id, task.serverless)
        self.drhouse_bucket = drhouse_s3_bucket(task.region)
        self.drhouse_prefix = drhouse_s3_prefix(task.cluster_id)
        self.drhouse_local_path = task.cache_prefix.rstrip("/")
        self.interval_secs = 30
        if task.provider == "aws":
            self.s3 = S3()
        elif task.provider == "gcp":
            self.s3 = GCS()
        elif task.provider == "azure":
            self.s3 = AzureBlob()
        else:
            raise ValueError("invalid provider: " + task.provider)

    def is_timing(self, ts: int) -> bool:
        return ts % self.interval_secs == 0

    def ready(self, ts: int) -> bool:
        return True

    def start(self, ts: int):
        now = int(time.time())
        if ts + 3600 < now:
            # supplement mode
            if ts % 3600 == 0:
                self._run(ts)
        else:
            # realtime mode
            hour_ts = ts - ts % 3600
            if ts % 3600 <= 100:
                self._run(hour_ts - 3600)
            self._run(hour_ts)

    def close(self, ts: int):
        pass

    def _run(self, ts: int):
        hour_ts = ts - ts % 3600
        utc_hour = datetime.utcfromtimestamp(hour_ts).strftime("%Y%m%d%H")

        # Download existed merged logs if needed.
        merged_regions_s3_key = f"{self.drhouse_prefix}/regions-merged/1h/{utc_hour}"
        merged_regions_local_path = f"{self.drhouse_local_path}/regions-merged/1h/{utc_hour}"
        log.info(f"[{self.cluster_id}] merged regions s3 key: {merged_regions_s3_key}")
        log.info(f"[{self.cluster_id}] merged regions local path: {merged_regions_local_path}")
        if not os.path.exists(merged_regions_local_path):
            log.info(f"[{self.cluster_id}] merged regions does not exists, downloading from s3")
            os.makedirs(os.path.dirname(merged_regions_local_path), exist_ok=True)
            ok = self.s3.download(self.drhouse_bucket, merged_regions_s3_key, merged_regions_local_path)
            if ok:
                log.info(f"[{self.cluster_id}] merged regions downloaded")
            else:
                log.info(f"[{self.cluster_id}] merged regions s3 object does not exists, will create empty one")
        else:
            log.info(f"[{self.cluster_id}] merged regions existed")

        try:
            with open(merged_regions_local_path, "a") as merged_regions_file:
                matched_keys = self.s3.list(self.o11y_bucket, f"{self.o11y_prefix}/regions/{utc_hour}")
                existed_keys = read_keys_from_merged_logs(merged_regions_local_path)
                added_keys = set(matched_keys) - set(existed_keys)
                log.info(f"[{self.cluster_id}] matched regions keys: {len(matched_keys)}")
                log.info(f"[{self.cluster_id}] existed regions keys: {len(existed_keys)}")
                log.info(f"[{self.cluster_id}] added regions keys: {len(added_keys)}")
                if not added_keys:
                    return
                objs = self.s3.objects(self.o11y_bucket, added_keys)
                for key, obj in zip(added_keys, objs):
                    try:
                        merged_regions_file.write(f"# {key}\n")
                        if key.endswith(".json.log.gz"):
                            key = key.removesuffix(".json.log.gz")
                        elif key.endswith(".json.log"):
                            key = key.removesuffix(".json.log")
                        utc_minute = key[len(key) - 12 :]
                        try:
                            content = read_content_with_retry(obj)
                            merged_regions_file.write(utc_minute)
                            merged_regions_file.write(content.decode())
                            merged_regions_file.write("\n")
                        except Exception:
                            log.exception("failed to read content with retry")
                    finally:
                        obj.close()
            self.s3.upload(self.drhouse_bucket, merged_regions_s3_key, merged_regions_local_path)
        finally:
            os.remove(merged_regions_local_path)


def o11y_s3_bucket(env: str, region: str, serverless: bool) -> str:
    if serverless:
        return f"o11y-prod-shared-{region}"
    if env.lower() == "prod":
        return f"o11y-prod-shared-{region}"
    return f"o11y-test-shared-{region}"


def o11y_s3_prefix(tenant: str, project: str, cluster: str, serverless: bool) -> str:
    if serverless:
        return f"0/prod-devtier/serverless/{tenant}/{project}/{cluster}"
    return f"0/{tenant}/{project}/{cluster}"


def drhouse_s3_bucket(region: str) -> str:
    return f"o11y-prod-shared-{region}"


def drhouse_s3_prefix(cluster: str) -> str:
    return f"diagnosis/data/{cluster}"


def read_content_with_retry(obj, max_retries=3):
    for attempt in range(max_retries):
        try:
            return obj.read()
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            log.warning(f"failed to read content (attempt {attempt + 1}/{max_retries}): {e}")
            time.sleep(1)