#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""Database table creation and insertion utilities."""

import time
from misc import _log


def _migrate_latency_columns(sqlite):
    """Add latency columns to existing databases."""
    try:
        sqlite.execute('SELECT avg_latency FROM proxylist LIMIT 1')
    except Exception:
        sqlite.execute('ALTER TABLE proxylist ADD COLUMN avg_latency REAL DEFAULT 0')
        sqlite.execute('ALTER TABLE proxylist ADD COLUMN latency_samples INT DEFAULT 0')
        sqlite.commit()


def update_proxy_latency(sqlite, proxy, latency_ms):
    """Update rolling average latency for a proxy.

    Args:
        sqlite: Database connection
        proxy: Proxy address (ip:port)
        latency_ms: Response latency in milliseconds
    """
    row = sqlite.execute(
        'SELECT avg_latency, latency_samples FROM proxylist WHERE proxy=?',
        (proxy,)
    ).fetchone()

    if row:
        old_avg, samples = row[0] or 0, row[1] or 0
        # Exponential moving average, capped at 100 samples
        new_samples = min(samples + 1, 100)
        if samples == 0:
            new_avg = latency_ms
        else:
            # Weight recent samples more heavily
            alpha = 2.0 / (new_samples + 1)
            new_avg = alpha * latency_ms + (1 - alpha) * old_avg

        sqlite.execute(
            'UPDATE proxylist SET avg_latency=?, latency_samples=? WHERE proxy=?',
            (new_avg, new_samples, proxy)
        )


def create_table_if_not_exists(sqlite, dbname):
    """Create database table with indexes if it doesn't exist."""
    if dbname == 'proxylist':
        sqlite.execute("""CREATE TABLE IF NOT EXISTS proxylist (
            proxy BLOB UNIQUE,
            country BLOB,
            added INT,
            failed INT,
            tested INT,
            dronebl INT,
            proto TEXT,
            mitm INT,
            success_count INT,
            ip TEXT,
            port INT,
            consecutive_success INT,
            total_duration INT,
            avg_latency REAL DEFAULT 0,
            latency_samples INT DEFAULT 0)""")
        # Indexes for common query patterns
        sqlite.execute('CREATE INDEX IF NOT EXISTS idx_proxylist_failed ON proxylist(failed)')
        sqlite.execute('CREATE INDEX IF NOT EXISTS idx_proxylist_tested ON proxylist(tested)')
        sqlite.execute('CREATE INDEX IF NOT EXISTS idx_proxylist_proto ON proxylist(proto)')
        # Migration: add latency columns if missing
        _migrate_latency_columns(sqlite)

    elif dbname == 'uris':
        sqlite.execute("""CREATE TABLE IF NOT EXISTS uris (
            url TEXT UNIQUE,
            content_type TEXT,
            check_time INT,
            error INT,
            stale_count INT,
            retrievals INT,
            proxies_added INT,
            added INT)""")
        # Indexes for common query patterns
        sqlite.execute('CREATE INDEX IF NOT EXISTS idx_uris_error ON uris(error)')
        sqlite.execute('CREATE INDEX IF NOT EXISTS idx_uris_checktime ON uris(check_time)')

    sqlite.commit()


def insert_proxies(proxydb, proxies, url):
    """Insert new proxies into database."""
    if not proxies:
        return
    timestamp = int(time.time())
    rows = []
    for p in proxies:
        ip, port = p.split(':')
        rows.append((timestamp, p, ip, port, 3, 0, 0, 0, 0, 0))
    proxydb.executemany(
        'INSERT OR IGNORE INTO proxylist '
        '(added,proxy,ip,port,failed,tested,success_count,total_duration,mitm,consecutive_success) '
        'VALUES (?,?,?,?,?,?,?,?,?,?)',
        rows
    )
    proxydb.commit()
    _log('+%d proxy/ies from %s' % (len(proxies), url), 'added')


def insert_urls(urls, search, sqlite):
    """Insert new URLs into database."""
    if not urls:
        return
    timestamp = int(time.time())
    rows = [(timestamp, u, 0, 1, 0, 0, 0) for u in urls]
    sqlite.executemany(
        'INSERT OR IGNORE INTO uris '
        '(added,url,check_time,error,stale_count,retrievals,proxies_added) '
        'VALUES (?,?,?,?,?,?,?)',
        rows
    )
    sqlite.commit()
    _log('+%d url(s) from %s' % (len(urls), search), 'added')