from __future__ import annotations

import logging
import os
import sys
import time
from datetime import timezone, timedelta
from pathlib import Path
from zoneinfo import ZoneInfo

from omegaconf import OmegaConf

_THIS_DIR = Path(__file__).resolve().parent
_YAML_PATH = _THIS_DIR / "config.yaml"

def _load_cfg():
    return OmegaConf.load(_YAML_PATH)

cfg = _load_cfg()

_LOG_TZ_NAME = os.environ.get("QUINTUS_LOG_TZ", "Asia/Kolkata")
try:
    _LOG_TZ = ZoneInfo(_LOG_TZ_NAME)
except Exception:
    _LOG_TZ = timezone(timedelta(hours=5, minutes=30))
    _LOG_TZ_NAME = "Asia/Kolkata"

os.environ["TZ"] = _LOG_TZ_NAME
if hasattr(time, "tzset"):
    time.tzset()
_LOG_TZ_LABEL = "IST" if _LOG_TZ_NAME == "Asia/Kolkata" else _LOG_TZ_NAME

def _read_bool_env(name: str) -> bool | None:
    raw = os.environ.get(name)
    if raw is None:
        return None

    normalised = raw.strip().lower()
    if normalised in {"1", "true", "yes", "on"}:
        return True
    if normalised in {"0", "false", "no", "off"}:
        return False
    raise ValueError(
        f"Invalid boolean value for {name}: {raw!r}. "
        "Use 1/0, true/false, yes/no, or on/off."
    )

# Environment variable overrides used by the wrapper.
if os.environ.get("QUINTUS_TEACHER_MODEL"):
    cfg.model.teacher = os.environ["QUINTUS_TEACHER_MODEL"]
if os.environ.get("QUINTUS_TEACHER_REVISION"):
    cfg.model.teacher_revision = os.environ["QUINTUS_TEACHER_REVISION"]
if os.environ.get("QUINTUS_STUDENT_MODEL"):
    cfg.model.student = os.environ["QUINTUS_STUDENT_MODEL"]
if os.environ.get("QUINTUS_STUDENT_REVISION"):
    cfg.model.student_revision = os.environ["QUINTUS_STUDENT_REVISION"]
if os.environ.get("QUINTUS_TOKENIZER_MODEL"):
    cfg.model.tokenizer = os.environ["QUINTUS_TOKENIZER_MODEL"]
if os.environ.get("QUINTUS_TOKENIZER_REVISION"):
    cfg.model.tokenizer_revision = os.environ["QUINTUS_TOKENIZER_REVISION"]
if os.environ.get("QUINTUS_STUDENT_DIR"):
    cfg.paths.student_dir = os.environ["QUINTUS_STUDENT_DIR"]
if os.environ.get("QUINTUS_TOKENIZER_DIR"):
    cfg.paths.tokenizer_dir = os.environ["QUINTUS_TOKENIZER_DIR"]
if os.environ.get("NUM_SAMPLES"):
    cfg.data.num_samples = int(os.environ["NUM_SAMPLES"])
if os.environ.get("TRAIN_NUM_EPOCHS"):
    cfg.training.num_epochs = int(os.environ["TRAIN_NUM_EPOCHS"])
if os.environ.get("TRAIN_LEARNING_RATE"):
    cfg.training.learning_rate = float(os.environ["TRAIN_LEARNING_RATE"])
if os.environ.get("TRAIN_ALPHA"):
    cfg.training.alpha = float(os.environ["TRAIN_ALPHA"])
if os.environ.get("TRAIN_TEMPERATURE"):
    cfg.training.temperature = float(os.environ["TRAIN_TEMPERATURE"])
if os.environ.get("TRAIN_TOP_K"):
    cfg.training.top_k = int(os.environ["TRAIN_TOP_K"])
if os.environ.get("QUINTUS_ONLINE_KD_TOKEN_CHUNK_SIZE"):
    cfg.training.online_kd_token_chunk_size = int(os.environ["QUINTUS_ONLINE_KD_TOKEN_CHUNK_SIZE"])
if os.environ.get("TRAIN_MICRO_BATCH_SIZE"):
    cfg.training.micro_batch_size = int(os.environ["TRAIN_MICRO_BATCH_SIZE"])
if os.environ.get("TRAIN_GRAD_ACCUM_STEPS"):
    cfg.training.grad_accum_steps = int(os.environ["TRAIN_GRAD_ACCUM_STEPS"])
if os.environ.get("TRAIN_DATALOADER_WORKERS"):
    cfg.training.dataloader_workers = int(os.environ["TRAIN_DATALOADER_WORKERS"])
if os.environ.get("TRAIN_PREFETCH_FACTOR"):
    cfg.training.prefetch_factor = int(os.environ["TRAIN_PREFETCH_FACTOR"])
sequence_packing_override = _read_bool_env("QUINTUS_SEQUENCE_PACKING")
if sequence_packing_override is not None:
    cfg.training.sequence_packing.enabled = sequence_packing_override
if os.environ.get("QUINTUS_PACK_LENGTH"):
    cfg.training.sequence_packing.pack_length = int(os.environ["QUINTUS_PACK_LENGTH"])
compile_override = _read_bool_env("QUINTUS_COMPILE_MODEL")
if compile_override is not None:
    cfg.training.compile_model = compile_override
fused_adamw_override = _read_bool_env("TRAIN_FUSED_ADAMW")
if fused_adamw_override is not None:
    cfg.training.fused_adamw = fused_adamw_override
if os.environ.get("QUINTUS_DISTILLED_DIR"):
    cfg.paths.distilled_dir = os.environ["QUINTUS_DISTILLED_DIR"]
if os.environ.get("DATA_STREAM_SHUFFLE_BUFFER_SIZE"):
    cfg.data.stream_shuffle_buffer_size = int(os.environ["DATA_STREAM_SHUFFLE_BUFFER_SIZE"])
if os.environ.get("DATA_STREAM_SHUFFLE_SEED"):
    cfg.data.stream_shuffle_seed = int(os.environ["DATA_STREAM_SHUFFLE_SEED"])
remote_code_override = _read_bool_env("QUINTUS_ALLOW_REMOTE_CODE")
if remote_code_override is not None:
    cfg.model.allow_remote_code = remote_code_override

class _TagFormatter(logging.Formatter):
    def __init__(self, tag: str, fmt: str, datefmt: str | None = None):
        super().__init__(fmt=fmt, datefmt=datefmt)
        self.tag = tag

    def formatTime(self, record: logging.LogRecord, datefmt: str | None = None) -> str:
        dt = datetime_from_timestamp(record.created)
        if datefmt:
            return dt.strftime(datefmt)
        return dt.isoformat(timespec="seconds")

    def format(self, record: logging.LogRecord) -> str:
        record.tag = self.tag  # type: ignore[attr-defined]
        return super().format(record)

def datetime_from_timestamp(timestamp: float):
    from datetime import datetime

    return datetime.fromtimestamp(timestamp, tz=_LOG_TZ)

def setup_logger(module_tag: str, rank: int = -1) -> logging.Logger:
    name = f"quintus.{module_tag}"
    logger = logging.getLogger(name)

    if logger.handlers:
        return logger

    logger.setLevel(logging.DEBUG)
    logger.propagate = False

    # Suppress duplicate output from non-primary ranks.
    if rank not in (-1, 0):
        logger.addHandler(logging.NullHandler())
        return logger

    # Plain text file handler.
    file_fmt = _TagFormatter(
        tag=module_tag,
        fmt=f"[%(asctime)s {_LOG_TZ_LABEL}] [%(levelname)-5s] [%(tag)-8s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    log_dir = os.path.dirname(cfg.paths.log_file)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    file_handler = logging.FileHandler(cfg.paths.log_file, mode="a", encoding="utf-8")
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(file_fmt)
    logger.addHandler(file_handler)

    # Plain text console handler. Keep the runtime logs stable across terminals,
    # notebooks and log files.
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(file_fmt)
    logger.addHandler(console_handler)

    return logger

def emit_log_spacing(logger: logging.Logger, count: int = 2) -> None:
    if count <= 0:
        return

    blank_block = "\n" * count
    for handler in logger.handlers:
        if isinstance(handler, logging.NullHandler):
            continue

        stream = getattr(handler, "stream", None)
        if stream is not None and hasattr(stream, "write"):
            stream.write(blank_block)
            flush = getattr(stream, "flush", None)
            if callable(flush):
                flush()
            continue

        console = getattr(handler, "console", None)
        if console is not None:
            console.print(blank_block, end="")