from __future__ import annotations

import argparse
import asyncio
import json
import math
import os
import statistics
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, AsyncIterator

from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles


DEFAULT_LOG = "/var/log/nginx-traffic/traffic.log"
MAX_HOSTS = 1000


@dataclass
class Bucket:
    incoming: int = 0
    outgoing: int = 0
    requests: int = 0
    statuses: list[int] = field(default_factory=lambda: [0, 0, 0, 0])
    request_times: list[float] = field(default_factory=list)
    upstream_times: list[float] = field(default_factory=list)


class Series:
    def __init__(self, retention: int):
        self.retention = retention
        self.buckets: dict[int, Bucket] = {}

    def bucket(self, second: int) -> Bucket:
        if second not in self.buckets:
            self.buckets[second] = Bucket()
        cutoff = second - self.retention + 1
        for old in [key for key in self.buckets if key < cutoff]:
            del self.buckets[old]
        return self.buckets[second]


def percentile(values: list[float], p: float) -> float | None:
    if not values:
        return None
    values = sorted(values)
    index = (len(values) - 1) * p
    low, high = math.floor(index), math.ceil(index)
    if low == high:
        return values[low]
    return values[low] + (values[high] - values[low]) * (index - low)


def parse_line(line: str) -> dict[str, Any] | None:
    """Parse the configured, whitespace-delimited Nginx traffic format."""
    fields = line.strip().split()
    # A trailing empty $upstream_response_time disappears when splitting on
    # whitespace, so seven fields is valid and means "no upstream timing".
    if len(fields) < 7:
        return None
    try:
        timestamp = float(fields[0])
        status = int(fields[3])
        request_length = int(fields[4])
        bytes_sent = int(fields[5])
        request_time = float(fields[6])
        upstream_raw = "".join(fields[7:]) if len(fields) > 7 else "-"
        upstream_values = [float(value) for value in upstream_raw.split(",") if value not in ("", "-")]
        if not all(math.isfinite(value) and value >= 0 for value in
                   (timestamp, request_length, bytes_sent, request_time, *upstream_values)):
            return None
        if status < 100 or status > 599:
            return None
    except (ValueError, OverflowError):
        return None
    return {
        "timestamp": timestamp,
        "host": fields[1] or "(unknown)",
        "incoming": request_length,
        "outgoing": bytes_sent,
        "status": status,
        "request_time": request_time,
        "upstream_time": statistics.fmean(upstream_values) if upstream_values else None,
    }


class Telemetry:
    def __init__(self, retention: int):
        self.retention = max(60, min(retention, 86400))
        self.global_series = Series(self.retention)
        self.hosts: dict[str, Series] = {}
        self.last_seen: dict[str, int] = {}
        self.total_incoming = 0
        self.total_outgoing = 0
        self.lock = threading.RLock()
        self.subscribers: set[asyncio.Queue[str]] = set()

    def ingest(self, event: dict[str, Any]) -> None:
        now = int(time.time())
        second = int(event["timestamp"])
        # Log timestamps are normally current, but clamping protects bounded
        # storage if a clock-jumped or otherwise surprising line appears.
        if second < now - self.retention:
            return
        second = min(second, now)
        with self.lock:
            self._add(self.global_series.bucket(second), event)
            host = event["host"]
            if host not in self.hosts:
                if len(self.hosts) >= MAX_HOSTS:
                    oldest = min(self.last_seen, key=self.last_seen.get)
                    self.hosts.pop(oldest, None)
                    self.last_seen.pop(oldest, None)
                self.hosts[host] = Series(self.retention)
            self._add(self.hosts[host].bucket(second), event)
            self.last_seen[host] = second
            self.total_incoming += event["incoming"]
            self.total_outgoing += event["outgoing"]

    @staticmethod
    def _add(bucket: Bucket, event: dict[str, Any]) -> None:
        bucket.incoming += event["incoming"]
        bucket.outgoing += event["outgoing"]
        bucket.requests += 1
        status_class = event["status"] // 100
        if 2 <= status_class <= 5:
            bucket.statuses[status_class - 2] += 1
        if len(bucket.request_times) < 256:
            bucket.request_times.append(event["request_time"])
        if event["upstream_time"] is not None and len(bucket.upstream_times) < 256:
            bucket.upstream_times.append(event["upstream_time"])

    def subscribe(self) -> asyncio.Queue[str]:
        queue: asyncio.Queue[str] = asyncio.Queue(maxsize=2)
        with self.lock:
            self.subscribers.add(queue)
        return queue

    def unsubscribe(self, queue: asyncio.Queue[str]) -> None:
        with self.lock:
            self.subscribers.discard(queue)

    def snapshot(self, window: int) -> dict[str, Any]:
        now = int(time.time())
        start = now - window + 1
        with self.lock:
            points = [self._point(self.global_series.buckets.get(second), second) for second in range(start, now + 1)]
            hosts = sorted(self.hosts, key=lambda host: self.last_seen.get(host, 0), reverse=True)
            active = [host for host in hosts if self.last_seen.get(host, 0) >= now - 59]
            host_points = {
                host: [self._point(self.hosts[host].buckets.get(second), second) for second in range(start, now + 1)]
                for host in hosts
            }
            current = self._summary(points[-60:])
            current["total_incoming"] = self.total_incoming
            current["total_outgoing"] = self.total_outgoing
            current["active_hosts"] = len(active)
            current["busiest_host"] = max(active, key=lambda host: sum(p["incoming"] + p["outgoing"] for p in host_points[host][-60:]), default=None)
            return {"type": "telemetry", "now": now, "points": points, "hosts": hosts, "host_points": host_points, "summary": current}

    @staticmethod
    def _point(bucket: Bucket | None, second: int) -> dict[str, Any]:
        if not bucket:
            return {"time": second, "incoming": 0, "outgoing": 0, "requests": 0, "status": [0, 0, 0, 0], "latency": None, "p95": None, "upstream": None}
        return {"time": second, "incoming": bucket.incoming, "outgoing": bucket.outgoing, "requests": bucket.requests, "status": bucket.statuses, "latency": statistics.fmean(bucket.request_times) if bucket.request_times else None, "p95": percentile(bucket.request_times, .95), "upstream": statistics.fmean(bucket.upstream_times) if bucket.upstream_times else None}

    @staticmethod
    def _summary(points: list[dict[str, Any]]) -> dict[str, Any]:
        seconds = max(1, len(points))
        total_requests = sum(point["requests"] for point in points)
        statuses = [sum(point["status"][index] for point in points) for index in range(4)]
        latencies = [point["latency"] for point in points if point["latency"] is not None]
        upstream = [point["upstream"] for point in points if point["upstream"] is not None]
        return {"incoming": sum(point["incoming"] for point in points) / seconds, "outgoing": sum(point["outgoing"] for point in points) / seconds, "requests": total_requests / seconds, "status": statuses, "latency": statistics.fmean(latencies) if latencies else None, "p95": percentile(latencies, .95), "upstream": statistics.fmean(upstream) if upstream else None}


async def follow_log(path: str, telemetry: Telemetry, stop: asyncio.Event) -> None:
    handle = None
    inode = None
    pending = ""
    while not stop.is_set():
        try:
            stat = await asyncio.to_thread(os.stat, path)
            if handle is None:
                handle = await asyncio.to_thread(open, path, "r", encoding="utf-8", errors="replace")
                await asyncio.to_thread(handle.seek, 0, os.SEEK_END)
                inode = (stat.st_dev, stat.st_ino)
            elif inode != (stat.st_dev, stat.st_ino):
                # Drain what is currently available from the renamed file before switching.
                for line in handle:
                    event = parse_line(line)
                    if event:
                        telemetry.ingest(event)
                handle.close()
                handle = await asyncio.to_thread(open, path, "r", encoding="utf-8", errors="replace")
                inode = (stat.st_dev, stat.st_ino)
                # Never join a partial line from the old inode to a new file.
                pending = ""
            data = await asyncio.to_thread(handle.read)
            if data:
                pending += data
                lines = pending.split("\n")
                pending = lines.pop()
                for line in lines:
                    event = parse_line(line)
                    if event:
                        telemetry.ingest(event)
            await asyncio.sleep(.15)
        except FileNotFoundError:
            if handle:
                data = handle.read()
                if data:
                    pending += data
            await asyncio.sleep(.5)
        except (OSError, ValueError):
            if handle:
                handle.close()
            handle, inode = None, None
            await asyncio.sleep(.5)
    if handle:
        handle.close()


def create_app(log_path: str = DEFAULT_LOG, retention: int = 3600) -> FastAPI:
    telemetry = Telemetry(retention)
    stop = asyncio.Event()
    app = FastAPI(title="Nginx Live Traffic Monitor")
    app.mount("/static", StaticFiles(directory=Path(__file__).parent / "static"), name="static")

    @app.on_event("startup")
    async def startup() -> None:
        app.state.task = asyncio.create_task(follow_log(log_path, telemetry, stop))
        app.state.publisher = asyncio.create_task(publish())

    @app.on_event("shutdown")
    async def shutdown() -> None:
        stop.set()
        for task in (getattr(app.state, "task", None), getattr(app.state, "publisher", None)):
            if task:
                task.cancel()

    async def publish() -> None:
        while not stop.is_set():
            payload = json.dumps(telemetry.snapshot(3600), separators=(",", ":"))
            with telemetry.lock:
                subscribers = list(telemetry.subscribers)
            for queue in subscribers:
                try:
                    queue.put_nowait(payload)
                except asyncio.QueueFull:
                    try:
                        queue.get_nowait()
                    except asyncio.QueueEmpty:
                        pass
                    queue.put_nowait(payload)
            await asyncio.sleep(1)

    @app.get("/")
    async def index() -> FileResponse:
        return FileResponse(Path(__file__).parent / "static" / "index.html")

    @app.get("/api/telemetry")
    async def telemetry_api(window: int = 300) -> dict[str, Any]:
        return telemetry.snapshot(max(60, min(window, 3600)))

    @app.get("/api/events")
    async def events(request: Request, window: int = 300) -> StreamingResponse:
        queue = telemetry.subscribe()
        initial = json.dumps(telemetry.snapshot(max(60, min(window, 3600))), separators=(",", ":"))
        async def stream() -> AsyncIterator[str]:
            try:
                yield f"data: {initial}\n\n"
                while True:
                    if await request.is_disconnected():
                        break
                    try:
                        payload = await asyncio.wait_for(queue.get(), timeout=15)
                        yield f"data: {payload}\n\n"
                    except asyncio.TimeoutError:
                        yield ": keep-alive\n\n"
            finally:
                telemetry.unsubscribe(queue)
        return StreamingResponse(stream(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})

    app.state.telemetry = telemetry
    return app


def main() -> None:
    parser = argparse.ArgumentParser(description="Live Nginx traffic dashboard")
    parser.add_argument("--log-path", default=os.getenv("TRAFFIC_LOG_PATH", DEFAULT_LOG))
    parser.add_argument("--host", default=os.getenv("LISTEN_ADDRESS", "127.0.0.1"))
    parser.add_argument("--port", type=int, default=int(os.getenv("LISTEN_PORT", "8080")))
    parser.add_argument("--retention", type=int, default=int(os.getenv("TELEMETRY_RETENTION", "3600")))
    args = parser.parse_args()
    import uvicorn
    uvicorn.run(create_app(args.log_path, args.retention), host=args.host, port=args.port)


if __name__ == "__main__":
    main()