from __future__ import annotations

import logging
import os
import time
from typing import Any

from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl


def remove_old_logs(log_file: str | None, num_stages: int) -> None:
    try:
        if not log_file:
            return
        try:
            if os.path.exists(log_file):
                os.remove(log_file)
        except Exception:
            pass
        # Per-stage logs and stats
        for sid in range(num_stages):
            try:
                p = f"{log_file}.stage{sid}.log"
                if os.path.exists(p):
                    os.remove(p)
            except Exception:
                pass
            try:
                p = f"{log_file}.stage{sid}.stats.jsonl"
                if os.path.exists(p):
                    os.remove(p)
            except Exception:
                pass
        # Orchestrator stats files
        try:
            p = f"{log_file}.orchestrator.stats.jsonl"
            if os.path.exists(p):
                os.remove(p)
        except Exception:
            pass
        try:
            p = f"{log_file}.overall.stats.jsonl"
            if os.path.exists(p):
                os.remove(p)
        except Exception:
            pass
    except Exception:
        pass


def configure_orchestrator_logger(logger: logging.Logger, log_file: str | None) -> None:
    try:
        if not log_file:
            return
        has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
        if not has_file_handler:
            fh = logging.FileHandler(log_file)
            fh.setLevel(logging.DEBUG)
            fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"))
            logger.addHandler(fh)
            logger.setLevel(logging.DEBUG)
    except Exception:
        pass


def init_stats_paths(enable_stats: bool, log_file: str | None) -> tuple[str | None, str | None]:
    stats_file: str | None = None
    overall_file: str | None = None
    try:
        if enable_stats and log_file:
            stats_file = f"{log_file}.orchestrator.stats.jsonl"
            overall_file = f"{log_file}.overall.stats.jsonl"
    except Exception:
        stats_file = None
        overall_file = None
    return stats_file, overall_file


def _safe_append_jsonl(path: str | None, record: dict[str, Any]) -> None:
    if not path:
        return
    try:
        _append_jsonl(path, record)  # type: ignore[arg-type]
    except Exception:
        pass


def log_transfer_tx(
    stats_file: str | None,
    from_stage: int,
    to_stage: int,
    request_id: Any,
    size_bytes: int,
    tx_time_ms: float,
    used_shm: bool,
) -> None:
    _safe_append_jsonl(
        stats_file,
        {
            "type": "transfer_stats",
            "from_stage": from_stage,
            "to_stage": to_stage,
            "request_id": request_id,
            "size_bytes": int(size_bytes),
            "tx_time_ms": float(tx_time_ms),
            "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
            "used_shm": bool(used_shm),
        },
    )


def log_transfer_rx(
    stats_file: str | None,
    from_stage: int,
    to_stage: int,
    request_id: Any,
    rx_bytes: int,
    rx_decode_time_ms: float,
    in_flight_time_ms: float,
) -> None:
    _safe_append_jsonl(
        stats_file,
        {
            "type": "transfer_rx_stats",
            "from_stage": from_stage,
            "to_stage": to_stage,
            "request_id": request_id,
            "rx_bytes": int(rx_bytes),
            "rx_decode_time_ms": float(rx_decode_time_ms),
            "in_flight_time_ms": float(in_flight_time_ms),
            "rx_time_per_kb_ms": (
                (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0
            ),
        },
    )


def log_transfer_total(
    stats_file: str | None,
    from_stage: int,
    to_stage: int,
    request_id: Any,
    size_bytes: int,
    tx_time_ms: float,
    in_flight_time_ms: float,
    rx_decode_time_ms: float,
    total_time_ms: float,
) -> None:
    _safe_append_jsonl(
        stats_file,
        {
            "type": "transfer_total_stats",
            "from_stage": from_stage,
            "to_stage": to_stage,
            "request_id": request_id,
            "size_bytes": int(size_bytes),
            "tx_time_ms": float(tx_time_ms),
            "in_flight_time_ms": float(in_flight_time_ms),
            "rx_decode_time_ms": float(rx_decode_time_ms),
            "total_time_ms": float(total_time_ms),
            "total_time_per_kb_ms": (
                float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0
            ),
        },
    )


def log_orchestrator_e2e(
    stats_file: str | None,
    request_id: Any,
    final_stage_id: int,
    e2e_time_ms: float,
    num_tokens_out: int,
    e2e_time_per_token_ms: float,
) -> None:
    _safe_append_jsonl(
        stats_file,
        {
            "type": "orchestrator_request_e2e",
            "request_id": request_id,
            "final_stage_id": final_stage_id,
            "e2e_time_ms": float(e2e_time_ms),
            "num_tokens_out": int(num_tokens_out),
            "e2e_time_per_token_ms": float(e2e_time_per_token_ms),
        },
    )


def log_orchestrator_summary(stats_file: str | None, summary: dict[str, Any]) -> None:
    _safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})


def log_overall_summary(overall_stats_file: str | None, summary: dict[str, Any]) -> None:
    _safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})


def log_overall_record(overall_stats_file: str | None, record: dict[str, Any]) -> None:
    _safe_append_jsonl(overall_stats_file, record)


def log_stage_request_stats(
    stats_file: str | None,
    stage_id: int,
    request_id: Any,
    batch_size: int,
    num_tokens_out: int,
    stage_gen_time_ms: float,
    tokens_per_s: float,
    rx_transfer_bytes: int,
    rx_decode_time_ms: float,
    rx_mbps: float,
) -> None:
    _safe_append_jsonl(
        stats_file,
        {
            "type": "stage_request_stats",
            "stage_id": stage_id,
            "request_id": request_id,
            "batch_size": int(batch_size),
            "num_tokens_out": int(num_tokens_out),
            "stage_gen_time_ms": float(stage_gen_time_ms),
            "tokens_per_s": float(tokens_per_s),
            "rx_transfer_bytes": int(rx_transfer_bytes),
            "rx_decode_time_ms": float(rx_decode_time_ms),
            "rx_mbps": float(rx_mbps),
        },
    )


def log_stage_running_avg(
    stats_file: str | None,
    stage_id: int,
    total_tokens: int,
    total_gen_time_ms: float,
    avg_tokens_per_s: float,
) -> None:
    _safe_append_jsonl(
        stats_file,
        {
            "type": "stage_running_avg",
            "stage_id": stage_id,
            "total_tokens": int(total_tokens),
            "total_gen_time_ms": float(total_gen_time_ms),
            "avg_tokens_per_s": float(avg_tokens_per_s),
        },
    )


def log_stage_batch_stats(
    stats_file: str | None,
    stage_id: int,
    batch_size: int,
    batch_gen_time_ms: float,
    request_ids: list[Any],
) -> None:
    _safe_append_jsonl(
        stats_file,
        {
            "type": "stage_batch_stats",
            "stage_id": stage_id,
            "batch_size": int(batch_size),
            "batch_gen_time_ms": float(batch_gen_time_ms),
            "request_ids": list(request_ids),
        },
    )


def compute_and_log_stage_request_stats(
    stats_file: str | None,
    stage_id: int,
    request_id: Any,
    batch_size: int,
    engine_outputs: list[Any],
    stage_gen_time_ms: float,
    rx_transfer_bytes: int,
    rx_decode_time_ms: float,
) -> None:
    """Compute per-request metrics and log them in one call."""
    num_tokens = count_tokens_from_outputs(engine_outputs)
    tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
    rx_mbps = (
        (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0)
        if rx_transfer_bytes > 0
        else 0.0
    )
    log_stage_request_stats(
        stats_file,
        stage_id,
        request_id,
        int(batch_size),
        int(num_tokens),
        float(stage_gen_time_ms),
        float(tokens_per_s),
        int(rx_transfer_bytes),
        float(rx_decode_time_ms),
        float(rx_mbps),
    )


# ----------------- Aggregation helpers for orchestrator -----------------


def record_stage_metrics(
    per_request: dict[int, dict[str, Any]],
    stage_req_counts: list[int],
    stage_total_time_ms: list[float],
    stage_total_tokens: list[int],
    stage_id: int,
    req_id: Any,
    metrics: dict[str, Any],
) -> None:
    try:
        stage_req_counts[stage_id] += 1
        stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
        rid_int = int(req_id)
        pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
        pr_stages = pr["stages"]  # type: ignore[index]
        pr_stages[stage_id] = {
            "stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)),
            "num_tokens_out": int(metrics.get("num_tokens_out", 0)),
        }
    except Exception:
        pass


def aggregate_rx_and_maybe_total(
    transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
    transfer_agg: dict[tuple[int, int], dict[str, float]],
    per_request: dict[int, dict[str, Any]],
    stage_id: int,
    req_id: Any,
    rx_bytes: float,
    rx_ms: float,
    in_flight_ms: float,
) -> tuple[int, float, float] | None:
    try:
        # Update RX aggregates for (stage_id-1 -> stage_id)
        if stage_id > 0:
            key = (stage_id - 1, stage_id)
            agg = transfer_agg.get(key)
            if agg is None:
                agg = {
                    "sum_bytes": 0.0,
                    "sum_ms": 0.0,
                    "count": 0.0,
                    "sum_rx_bytes": 0.0,
                    "sum_rx_ms": 0.0,
                    "rx_count": 0.0,
                    "sum_total_ms": 0.0,
                    "total_count": 0.0,
                }
                transfer_agg[key] = agg
            agg["sum_rx_bytes"] += float(rx_bytes)
            agg["sum_rx_ms"] += float(rx_ms)
            agg["rx_count"] += 1.0

            # Try combine with sender-side timing if present
            try:
                rid_int = int(req_id)
            except Exception:
                rid_int = req_id  # best effort
            s = transfer_edge_req.get((stage_id - 1, stage_id, rid_int))
            if s is None:
                return None
            tx_ms = float(s.get("tx_ms", 0.0))
            size_b = float(s.get("size_bytes", rx_bytes))
            total_ms = tx_ms + float(in_flight_ms) + float(rx_ms)
            agg["sum_total_ms"] += total_ms
            agg["total_count"] += 1.0
            # accumulate per-request transfer totals
            try:
                pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
                pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms  # type: ignore[index]
                pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes)  # type: ignore[index]
            except Exception:
                pass
            return int(size_b), float(tx_ms), float(total_ms)
        return None
    except Exception:
        return None


def record_sender_transfer_agg(
    transfer_agg: dict[tuple[int, int], dict[str, float]],
    transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
    from_stage: int,
    to_stage: int,
    req_id: Any,
    size_bytes: int,
    tx_ms: float,
) -> None:
    try:
        key = (from_stage, to_stage)
        agg = transfer_agg.get(key)
        if agg is None:
            agg = {
                "sum_bytes": 0.0,
                "sum_ms": 0.0,
                "count": 0.0,
                "sum_rx_bytes": 0.0,
                "sum_rx_ms": 0.0,
                "rx_count": 0.0,
                "sum_total_ms": 0.0,
                "total_count": 0.0,
            }
            transfer_agg[key] = agg
        agg["sum_bytes"] += float(size_bytes)
        agg["sum_ms"] += float(tx_ms)
        agg["count"] += 1.0
        # Store sender-side timing for per-request combination
        try:
            rid_int = int(req_id)
        except Exception:
            rid_int = req_id
        transfer_edge_req[(from_stage, to_stage, rid_int)] = {
            "tx_ms": float(tx_ms),
            "size_bytes": float(size_bytes),
        }
    except Exception:
        pass


def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
    total = 0
    for _ro in engine_outputs:
        try:
            outs = getattr(_ro, "outputs", None)
            if outs and len(outs) > 0:
                tokens = getattr(outs[0], "token_ids", None)
                if tokens is not None:
                    total += len(tokens)
        except Exception:
            pass
    return total


def build_stage_summary(
    stage_req_counts: list[int],
    stage_total_tokens: list[int],
    stage_total_time_ms: list[float],
) -> list[dict[str, Any]]:
    summary: list[dict[str, Any]] = []
    for sid in range(len(stage_req_counts)):
        reqs = stage_req_counts[sid]
        tokens = stage_total_tokens[sid]
        total_ms = float(stage_total_time_ms[sid])
        avg_req = (total_ms / reqs) if reqs > 0 else 0.0
        avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
        summary.append(
            {
                "stage_id": sid,
                "requests": int(reqs),
                "tokens": int(tokens),
                "total_time_ms": total_ms,
                "avg_time_per_request_ms": avg_req,
                "avg_tokens_per_s": avg_tok,
            }
        )
    return summary


def build_transfer_summary(
    transfer_agg: dict[tuple[int, int], dict[str, float]],
) -> list[dict[str, Any]]:
    summary: list[dict[str, Any]] = []
    for (src, dst), agg in transfer_agg.items():
        sum_bytes = float(agg.get("sum_bytes", 0.0))
        sum_ms = float(agg.get("sum_ms", 0.0))
        samples = int(agg.get("count", 0.0))
        tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
        sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0))
        sum_rx_ms = float(agg.get("sum_rx_ms", 0.0))
        samples_rx = int(agg.get("rx_count", 0.0))
        rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0
        sum_total_ms = float(agg.get("sum_total_ms", 0.0))
        samples_total = int(agg.get("total_count", 0.0))
        total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
        summary.append(
            {
                "from_stage": src,
                "to_stage": dst,
                "samples": samples,
                "total_bytes": int(sum_bytes),
                "total_time_ms": sum_ms,
                "tx_mbps": tx_mbps,
                "rx_samples": samples_rx,
                "rx_total_bytes": int(sum_rx_bytes),
                "rx_total_time_ms": sum_rx_ms,
                "rx_mbps": rx_mbps,
                "total_samples": samples_total,
                "total_transfer_time_ms": sum_total_ms,
                "total_mbps": total_mbps,
            }
        )
    return summary


class OrchestratorMetrics:
    def __init__(
        self,
        num_stages: int,
        enable_stats: bool,
        stats_file: str | None,
        overall_stats_file: str | None,
        wall_start_ts: float,
    ) -> None:
        self.num_stages = int(num_stages)
        self.enable_stats = bool(enable_stats)
        self.stats_file = stats_file
        self.overall_stats_file = overall_stats_file
        self.stage_total_time_ms: list[float] = [0.0 for _ in range(self.num_stages)]
        self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)]
        self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)]
        self.transfer_agg: dict[tuple[int, int], dict[str, float]] = {}
        self.transfer_edge_req: dict[tuple[int, int, int], dict[str, float]] = {}
        self.e2e_total_ms: float = 0.0
        self.e2e_total_tokens: int = 0
        self.e2e_count: int = 0
        self.e2e_done: set[int] = set()
        self.per_request: dict[int, dict[str, Any]] = {}
        self.sum_per_request_transfer_ms: float = 0.0
        self.wall_start_ts: float = float(wall_start_ts)
        self.last_finish_ts: float = float(wall_start_ts)
        self.stage_seen_batches: dict[int, set] = {sid: set() for sid in range(self.num_stages)}
        self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)]
        self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)]

    def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: dict[str, Any]) -> None:
        record_stage_metrics(
            self.per_request,
            self.stage_req_counts,
            self.stage_total_time_ms,
            self.stage_total_tokens,
            stage_id,
            req_id,
            metrics,
        )
        try:
            batch_id_raw = metrics.get("batch_id", None)
            if batch_id_raw is not None:
                batch_id = int(batch_id_raw)
                if batch_id not in self.stage_seen_batches[stage_id]:
                    self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0))
                    self.stage_seen_batches[stage_id].add(batch_id)
        except Exception:
            pass
        rx_b = float(metrics.get("rx_transfer_bytes", 0.0))
        rx_ms = float(metrics.get("rx_decode_time_ms", 0.0))
        in_flight_ms = float(metrics.get("rx_in_flight_time_ms", 0.0))
        combined = aggregate_rx_and_maybe_total(
            self.transfer_edge_req,
            self.transfer_agg,
            self.per_request,
            stage_id,
            req_id,
            rx_b,
            rx_ms,
            in_flight_ms,
        )
        if self.enable_stats and self.stats_file and stage_id > 0:
            log_transfer_rx(
                self.stats_file,
                stage_id - 1,
                stage_id,
                req_id,
                int(rx_b),
                float(rx_ms),
                float(in_flight_ms),
            )
            if combined is not None:
                size_b_c, tx_ms_c, total_ms_c = combined
                log_transfer_total(
                    self.stats_file,
                    stage_id - 1,
                    stage_id,
                    req_id,
                    int(size_b_c),
                    float(tx_ms_c),
                    float(in_flight_ms),
                    float(rx_ms),
                    float(total_ms_c),
                )

    def on_forward(
        self,
        from_stage: int,
        to_stage: int,
        req_id: Any,
        size_bytes: int,
        tx_ms: float,
        used_shm: bool,
    ) -> None:
        # Mark first input time for the destination stage if not set
        if self.stage_first_ts[to_stage] is None:
            self.stage_first_ts[to_stage] = time.time()
        if self.enable_stats and self.stats_file:
            log_transfer_tx(
                self.stats_file,
                from_stage,
                to_stage,
                req_id,
                int(size_bytes),
                float(tx_ms),
                bool(used_shm),
            )
        record_sender_transfer_agg(
            self.transfer_agg,
            self.transfer_edge_req,
            from_stage,
            to_stage,
            req_id,
            int(size_bytes),
            float(tx_ms),
        )

    def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: list[Any], req_start_ts: float) -> None:
        try:
            rid_int = int(req_id)
        except Exception:
            rid_int = req_id
        _t0 = float(req_start_ts)
        _t1 = time.time()
        # Update last output time for this stage
        prev_last = self.stage_last_ts[stage_id]
        self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1)
        self.last_finish_ts = max(self.last_finish_ts, _t1)
        e2e_ms = (_t1 - _t0) * 1000.0
        num_tokens = count_tokens_from_outputs(engine_outputs)
        self.e2e_total_ms += e2e_ms
        self.e2e_total_tokens += int(num_tokens)
        self.e2e_count += 1
        self.e2e_done.add(rid_int)
        pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
        per_req_record = {
            "type": "overall_request",
            "request_id": rid_int,
            "e2e_time_ms": e2e_ms,
            "num_tokens_out": int(num_tokens),
            "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
            "transfers_total_bytes": int(pr.get("transfers_bytes", 0)),
            "stages": pr.get("stages", {}),
        }
        self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
        if self.enable_stats and self.overall_stats_file:
            log_overall_record(self.overall_stats_file, per_req_record)
        if self.enable_stats and self.stats_file:
            e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
            log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)

    def build_and_log_summary(self, final_stage_id_for_e2e: int) -> dict[str, Any]:
        # Compute stage summary using wall time between first input and last output per stage
        stage_summary: list[dict[str, Any]] = []
        for sid in range(self.num_stages):
            first_ts = self.stage_first_ts[sid]
            last_ts = self.stage_last_ts[sid]
            total_ms = (
                (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
            )
            reqs = self.stage_req_counts[sid]
            tokens = self.stage_total_tokens[sid]
            avg_req = (total_ms / reqs) if reqs > 0 else 0.0
            avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0
            stage_summary.append(
                {
                    "stage_id": sid,
                    "requests": int(reqs),
                    "tokens": int(tokens),
                    "total_time_ms": float(total_ms),
                    "avg_time_per_request_ms": float(avg_req),
                    "avg_tokens_per_s": float(avg_tok),
                }
            )
        transfer_summary = build_transfer_summary(self.transfer_agg)
        e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
        e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0
        wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
        summary: dict[str, Any] = {
            "e2e_requests": int(self.e2e_count),
            "e2e_total_time_ms": float(wall_time_ms),
            "e2e_sum_time_ms": float(self.e2e_total_ms),
            "e2e_total_tokens": int(self.e2e_total_tokens),
            "e2e_avg_time_per_request_ms": e2e_avg_req,
            "e2e_avg_tokens_per_s": e2e_avg_tok,
            "wall_time_ms": wall_time_ms,
            "final_stage_id": final_stage_id_for_e2e,
            "stages": stage_summary,
            "transfers": transfer_summary,
        }
        if self.enable_stats and self.stats_file:
            log_orchestrator_summary(self.stats_file, summary)
        if self.enable_stats and self.overall_stats_file:
            log_overall_summary(self.overall_stats_file, summary)
        return summary