import asyncio
from aiohttp import web, WSMsgType

from datetime import datetime, timedelta, timezone
from collections import defaultdict
import html

def create_stats_structure():
    """Creates the nested dictionary structure for storing statistics."""
    def nested_dd():
        return defaultdict(lambda: defaultdict(int))
    return defaultdict(nested_dd)

def get_time_keys(dt: datetime):
    """Generates dictionary keys for different time granularities."""
    return {
        "hour": dt.strftime('%Y-%m-%d-%H'),
        "day": dt.strftime('%Y-%m-%d'),
        "week": dt.strftime('%Y-%W'), # Week number, Monday is first day
        "month": dt.strftime('%Y-%m'),
    }

def update_stats_counters(stats_dict: defaultdict, now: datetime):
    """Increments the appropriate time-based counters in a stats dictionary."""
    keys = get_time_keys(now)
    stats_dict['by_hour'][keys['hour']] += 1
    stats_dict['by_day'][keys['day']] += 1
    stats_dict['by_week'][keys['week']] += 1
    stats_dict['by_month'][keys['month']] += 1

def generate_time_series_svg(title: str, data: list[tuple[str, int]], y_label: str) -> str:
    """Generates a responsive SVG bar chart for time-series data."""
    if not data:
        return f"<h3>{html.escape(title)}</h3><p>No data yet.</p>"
    max_val = max(item[1] for item in data) if data else 1
    svg_height, svg_width = 250, 600
    bar_padding = 5
    bar_width = (svg_width - 50) / len(data) - bar_padding
    
    bars = ""
    labels = ""
    for i, (key, val) in enumerate(data):
        bar_height = (val / max_val) * (svg_height - 50) if max_val > 0 else 0
        x = i * (bar_width + bar_padding) + 40
        y = svg_height - bar_height - 30
        
        bars += f'<rect x="{x}" y="{y}" width="{bar_width}" height="{bar_height}" fill="#007BFF"><title>{html.escape(key)}: {val}</title></rect>'
        labels += f'<text x="{x + bar_width / 2}" y="{svg_height - 15}" font-size="11" text-anchor="middle">{html.escape(key)}</text>'

    return f"""
    <h3>{html.escape(title)}</h3>
    <div style="border:1px solid #ccc; padding: 10px; border-radius: 5px;">
        <svg viewBox="0 0 {svg_width} {svg_height}" style="width:100%; height:auto;">
            <g>{bars}</g>
            <g>{labels}</g>
            <line x1="35" y1="10" x2="35" y2="{svg_height - 30}" stroke="#aaa" stroke-width="1" />
            <line x1="35" y1="{svg_height - 30}" x2="{svg_width - 10}" y2="{svg_height - 30}" stroke="#aaa" stroke-width="1" />
            <text x="5" y="{svg_height - 30}" font-size="12">0</text>
            <text x="5" y="20" font-size="12">{max_val}</text>
        </svg>
    </div>
    """

@web.middleware
async def middleware(request, handler):
    """Middleware to count all incoming HTTP requests."""
    # Avoid counting requests to the stats page itself
    if request.path.startswith('/stats.html'):
        return await handler(request)
        
    update_stats_counters(request.app['stats']['http_requests'], datetime.now(timezone.utc))
    return await handler(request)

def update_websocket_stats(app):
    update_stats_counters(app['stats']['websocket_requests'], datetime.now(timezone.utc))

async def pipe_and_count_websocket(ws_from, ws_to, stats_dict):
    """This function proxies WebSocket messages AND counts them."""
    async for msg in ws_from:
        # This is the key part for monitoring WebSockets
        update_stats_counters(stats_dict, datetime.now(timezone.utc))
        
        if msg.type == WSMsgType.TEXT:
            await ws_to.send_str(msg.data)
        elif msg.type == WSMsgType.BINARY:
            await ws_to.send_bytes(msg.data)
        elif msg.type in (WSMsgType.CLOSE, WSMsgType.ERROR):
            await ws_to.close(code=ws_from.close_code)
            break


async def stats_handler(request: web.Request):
    """Handler to display the statistics dashboard."""
    stats = request.app['stats']
    now = datetime.now(timezone.utc)
    
    # Helper to prepare data for charts
    def get_data(source, period, count):
        data = []
        for i in range(count - 1, -1, -1):
            if period == 'hour':
                dt = now - timedelta(hours=i)
                key, label = dt.strftime('%Y-%m-%d-%H'), dt.strftime('%H:00')
                data.append((label, source['by_hour'].get(key, 0)))
            elif period == 'day':
                dt = now - timedelta(days=i)
                key, label = dt.strftime('%Y-%m-%d'), dt.strftime('%a')
                data.append((label, source['by_day'].get(key, 0)))
        return data

    http_hourly = get_data(stats['http_requests'], 'hour', 24)
    ws_hourly = get_data(stats['ws_messages'], 'hour', 24)
    http_daily = get_data(stats['http_requests'], 'day', 7)
    ws_daily = get_data(stats['ws_messages'], 'day', 7)

    body = f"""
    <html><head><title>App Stats</title><meta http-equiv="refresh" content="30"></head>
    <body>
        <h2>Application Dashboard</h2>
        <h3>Last 24 Hours</h3>
        {generate_time_series_svg("HTTP Requests", http_hourly, "Reqs/Hour")}
        {generate_time_series_svg("WebSocket Messages", ws_hourly, "Msgs/Hour")}
        <h3>Last 7 Days</h3>
        {generate_time_series_svg("HTTP Requests", http_daily, "Reqs/Day")}
        {generate_time_series_svg("WebSocket Messages", ws_daily, "Msgs/Day")}
    </body></html>
    """
    return web.Response(text=body, content_type='text/html')

