#!/usr/bin/env python3
"""
Shared monitoring core used by whp-log-monitor-ai.py and whp-monitor-poll.py.

Provides:
  - compute_signature(): deterministic SHA-256 signature for dedup
  - ten_minute_bucket(): 600s bucket of a unix timestamp
  - connect_mysql(): PDO-equivalent pymysql connection using /root/.my.cnf
  - connect_sqlite(): opens /var/lib/whp-ai-monitor/monitor.db
  - load_active_subscriptions(): returns list of subscribed primary domains
  - is_suppressed(): check if a signature is currently suppressed
  - insert_or_increment_issue(): idempotent issue insertion with dedup
  - call_ai(): basic/advanced model call via LiteLLM endpoint
"""
import hashlib
import json
import os
import sqlite3
import time
import configparser

try:
    import pymysql
except ImportError:
    pymysql = None  # type: ignore

try:
    import requests
except ImportError:
    requests = None  # type: ignore

AI_CONFIG_PATH = '/etc/whp/ai-monitor.conf'
AI_KEY_PATH = '/etc/whp/ai-monitor.key'
SQLITE_PATH = '/var/lib/whp-ai-monitor/monitor.db'
SETTINGS_JSON_PATH = '/docker/whp/settings.json'

# ---------- Signatures ----------

def ten_minute_bucket(ts: int) -> int:
    """Floor a unix timestamp to the nearest 600s."""
    return (ts // 600) * 600

def compute_signature(kind: str, **kwargs) -> str:
    """
    Deterministic signature for dedup.
      kind='health' + subsystem + severity  -> server-wide
      kind='site'   + site_id + category + bucket -> site-linked
    """
    if kind == 'health':
        raw = f"health:{kwargs['subsystem']}:{kwargs['severity']}"
    elif kind == 'site':
        raw = f"site:{kwargs['site_id']}:{kwargs['category']}:{kwargs['bucket']}"
    else:
        raise ValueError(f"unknown signature kind: {kind}")
    return hashlib.sha256(raw.encode('utf-8')).hexdigest()

# ---------- DB connections ----------

def _read_mysql_creds():
    """Parse /root/.my.cnf [client] section."""
    cp = configparser.ConfigParser()
    cp.read('/root/.my.cnf')
    s = cp['client']
    return {
        'host': s.get('host', '127.0.0.1'),
        'port': int(s.get('port', '3306')),
        'user': s.get('user'),
        'password': (s.get('password') or '').strip().strip('"').strip("'"),
    }

def connect_mysql():
    if pymysql is None:
        raise RuntimeError(
            "pymysql not installed; install with: pip install pymysql")
    creds = _read_mysql_creds()
    return pymysql.connect(
        host=creds['host'], port=creds['port'],
        user=creds['user'], password=creds['password'],
        database='whp', charset='utf8mb4', autocommit=False,
        cursorclass=pymysql.cursors.DictCursor,
    )

def connect_sqlite():
    conn = sqlite3.connect(SQLITE_PATH)
    conn.row_factory = sqlite3.Row
    return conn

# ---------- Settings ----------

def customer_monitoring_enabled() -> bool:
    """Read global toggle from /docker/whp/settings.json."""
    try:
        with open(SETTINGS_JSON_PATH) as f:
            data = json.load(f)
        cm = data.get('customer_monitoring') or {}
        return bool(cm.get('enabled'))
    except (FileNotFoundError, json.JSONDecodeError):
        return False

def load_active_subscriptions(mysql_conn) -> list:
    """Return list of dicts with keys {primary_domain, site_id, username}
    for all ai_monitor_subscriptions in status='active'."""
    with mysql_conn.cursor() as cur:
        cur.execute("""
            SELECT DISTINCT d.domain_name AS primary_domain, s.id AS site_id, s.username
            FROM whp.ai_monitor_subscriptions sub
            JOIN whp.sites s ON sub.site_id = s.id
            JOIN whp.domains d ON s.primary_domain_id = d.id
            WHERE sub.status = 'active'
        """)
        return cur.fetchall()

# ---------- Suppressions ----------

def is_suppressed(mysql_conn, site_id, signature) -> bool:
    """Check active suppression (site_id may be None for server-wide)."""
    with mysql_conn.cursor() as cur:
        if site_id is None:
            cur.execute(
                "SELECT id FROM whp.ai_monitor_suppressions "
                "WHERE signature=%s AND site_id IS NULL LIMIT 1",
                (signature,))
        else:
            cur.execute(
                "SELECT id FROM whp.ai_monitor_suppressions "
                "WHERE signature=%s AND site_id=%s LIMIT 1",
                (signature, site_id))
        return cur.fetchone() is not None

def bump_suppression_counter(mysql_conn, site_id, signature):
    with mysql_conn.cursor() as cur:
        if site_id is None:
            cur.execute(
                "UPDATE whp.ai_monitor_suppressions "
                "SET suppressed_occurrence_count = suppressed_occurrence_count + 1, "
                "    last_suppressed_event_at = NOW() "
                "WHERE signature=%s AND site_id IS NULL",
                (signature,))
        else:
            cur.execute(
                "UPDATE whp.ai_monitor_suppressions "
                "SET suppressed_occurrence_count = suppressed_occurrence_count + 1, "
                "    last_suppressed_event_at = NOW() "
                "WHERE signature=%s AND site_id=%s",
                (signature, site_id))
    mysql_conn.commit()

# ---------- Issue insertion ----------

def insert_or_increment_issue(mysql_conn, *, signature, site_id, subscription_id,
                              username, container_name,
                              severity, category, subsystem, source,
                              title, diagnosis, recommendation,
                              raw_log_excerpt, context_data):
    """
    Idempotent Issue write. Owns its own transaction — commits on success,
    rolls back on any exception.

      - If suppressed: bump counter, return (None, 'suppressed').
      - If open issue exists with same signature: increment occurrence_count,
        append source to sources JSON, return (id, 'incremented').
      - Otherwise: insert new row, return (id, 'created').

    Note: concurrent pollers can race on the SELECT → INSERT path. Follow-up
    will add a generated-column unique index or advisory lock. For now, an
    accidental duplicate is acceptable (two open rows; can be merged manually).
    """
    if is_suppressed(mysql_conn, site_id, signature):
        bump_suppression_counter(mysql_conn, site_id, signature)
        return (None, 'suppressed')

    try:
        with mysql_conn.cursor() as cur:
            cur.execute(
                "SELECT id, sources FROM whp.ai_monitor_issues "
                "WHERE signature=%s AND status='open' LIMIT 1",
                (signature,))
            row = cur.fetchone()

            if row:
                existing_sources = json.loads(row['sources'] or '[]')
                if source not in existing_sources:
                    existing_sources.append(source)
                cur.execute(
                    "UPDATE whp.ai_monitor_issues "
                    "SET occurrence_count = occurrence_count + 1, "
                    "    last_seen = NOW(), "
                    "    sources = %s "
                    "WHERE id=%s",
                    (json.dumps(existing_sources), row['id']))
                mysql_conn.commit()
                return (row['id'], 'incremented')

            cur.execute(
                """INSERT INTO whp.ai_monitor_issues
                    (subscription_id, site_id, username, container_name,
                     severity, category, subsystem,
                     source, sources, title, diagnosis, recommendation,
                     raw_log_excerpt, context_data, signature,
                     occurrence_count, first_seen, last_seen,
                     status, email_sent)
                   VALUES (%s, %s, %s, %s,
                           %s, %s, %s,
                           %s, %s, %s, %s, %s,
                           %s, %s, %s,
                           1, NOW(), NOW(),
                           'open', 0)""",
                (subscription_id, site_id, username, container_name,
                 severity, category, subsystem,
                 source, json.dumps([source]), title, diagnosis, recommendation,
                 raw_log_excerpt, json.dumps(context_data) if context_data else None,
                 signature))
            new_id = cur.lastrowid
        mysql_conn.commit()
        return (new_id, 'created')
    except Exception:
        try:
            mysql_conn.rollback()
        except Exception:
            pass
        raise

def mark_auto_resolved(mysql_conn, signature):
    """
    System-observed recovery: close the matching open issue.
    Sets status='fixed', auto_resolved_at=NOW(), status_changed_by='system',
    status_changed_at=NOW() on the open row with this signature.
    """
    with mysql_conn.cursor() as cur:
        cur.execute(
            "UPDATE whp.ai_monitor_issues "
            "SET status = 'fixed', "
            "    auto_resolved_at = COALESCE(auto_resolved_at, NOW()), "
            "    status_changed_by = 'system', "
            "    status_changed_at = NOW() "
            "WHERE signature=%s AND status='open'",
            (signature,))
    mysql_conn.commit()

# ---------- AI ----------

def _read_ai_config():
    cp = configparser.ConfigParser()
    cp.read(AI_CONFIG_PATH)
    return {
        'base_url': cp.get('ai', 'base_url', fallback='http://localhost:4000/v1'),
        'basic_model': cp.get('ai', 'basic_model',
                              fallback='global.anthropic.claude-haiku-4-5-20251001-v1:0'),
        'advanced_model': cp.get('ai', 'advanced_model',
                                 fallback='global.anthropic.claude-sonnet-4-6'),
    }

def _read_ai_key():
    with open(AI_KEY_PATH) as f:
        return f.read().strip()

def call_ai(prompt: str, tier: str = 'basic', max_tokens: int = 800, timeout: int = 30) -> dict:
    """
    tier: 'basic' or 'advanced'. Returns {'content': str, 'model': str, 'tokens': int}.
    Raises requests.exceptions on HTTP errors.
    """
    if requests is None:
        raise RuntimeError(
            "requests not installed; install with: pip install requests")
    cfg = _read_ai_config()
    model = cfg['basic_model'] if tier == 'basic' else cfg['advanced_model']
    resp = requests.post(
        f"{cfg['base_url'].rstrip('/')}/chat/completions",
        headers={'Authorization': f'Bearer {_read_ai_key()}',
                 'Content-Type': 'application/json'},
        json={
            'model': model,
            'messages': [{'role': 'user', 'content': prompt}],
            'max_tokens': max_tokens,
        },
        timeout=timeout,
    )
    resp.raise_for_status()
    j = resp.json()
    return {
        'content': j['choices'][0]['message']['content'],
        'model': model,
        'tokens': j.get('usage', {}).get('total_tokens', 0),
    }


# ---------- Phase-2 remediation guards ----------
#
# These helpers are shared between the cron-driven log monitor (which acts
# on log-pattern detections) and the minute-cadence health poll (which acts
# on health-API state transitions). Both call attempt_remediation() so the
# guards (cooldown, backup-protection, confidence gate, dry-run, allowlist)
# are enforced uniformly.

# Maps each issue category to a single safe action. Hardcoded here -- root
# can see the mapping on the AI Monitor settings page; making it user-
# editable was deferred since the consequences of a wrong mapping are
# severe. To override, edit this constant.
#
# Action grammar:
#   ('none',)                          -- never auto-remediate this category
#   ('container', '<name>')            -- docker restart <name>
#   ('container', 'auto')              -- docker restart on the AI-named target
#   ('service',   '<name>')            -- systemctl restart <name>
CATEGORY_ACTIONS = {
    'haproxy_unhealthy':  ('container', 'haproxy-manager'),
    'connection_storm':   ('container', 'haproxy-manager'),
    'docker_degraded':    ('container', 'auto'),
    'swap_pressure':      ('container', 'auto'),
    # CPU pressure: restart the AI-named container. Safe because:
    # - require_high_confidence gate ensures the AI is sure the named
    #   container is the cause AND a restart will fix it
    # - cooldown prevents restart-loops
    # - the per-target allowlist + customer-site lookup blocks unrelated
    #   container restarts
    # - backup-protection guard refuses while a backup is running
    'cpu_pressure':       ('container', 'auto'),
    'disk_full':          ('none',),  # cleanup is click-to-execute, not automatic
}


def category_action(category: str):
    """Return the configured action tuple for a category, or ('none',)."""
    return CATEGORY_ACTIONS.get(category, ('none',))


def is_backup_running(mysql_conn) -> bool:
    """True if any backup_queue row is currently in flight. The remediation
    engine refuses to restart anything when this is true unless the system
    is at imminent risk (see system_at_risk)."""
    try:
        with mysql_conn.cursor() as cur:
            cur.execute(
                "SELECT 1 FROM whp.backup_queue "
                "WHERE status IN ('queued','processing','running') LIMIT 1"
            )
            return cur.fetchone() is not None
    except Exception:
        # Fail-safe: if we can't tell, assume yes -- better to no-op than
        # interrupt a backup we couldn't see.
        return True


def system_at_risk(load_multiplier: float = 4.0,
                   swap_pct_threshold: float = 90.0) -> bool:
    """Return True only when the system is at imminent risk of collapse:
    1-min loadavg > load_multiplier x cpu_count AND swap usage > threshold.
    This is the only condition under which the remediation engine is
    allowed to override the backup-protection guard."""
    try:
        with open('/proc/loadavg') as f:
            load1 = float(f.read().split()[0])
        cores = os.cpu_count() or 1
        if load1 <= load_multiplier * cores:
            return False
        with open('/proc/meminfo') as f:
            mem = {k.strip(): v.strip() for k, v in
                   (line.split(':', 1) for line in f if ':' in line)}
        st = int(mem.get('SwapTotal', '0 kB').split()[0])
        sf = int(mem.get('SwapFree',  '0 kB').split()[0])
        if st <= 0:
            return False
        used_pct = 100.0 * (st - sf) / st
        return used_pct >= swap_pct_threshold
    except Exception:
        return False


def _is_whp_customer_container(mysql_conn, name: str) -> bool:
    """True when `name` is a docker container belonging to a WHP-managed
    customer site (named `<safe-domain>-NN` per setup convention).
    Lets the remediation engine restart customer site containers without
    enumerating each one in the static allowlist.
    """
    if not name:
        return False
    # First a cheap structural check: customer container names always end
    # with `-NN` (zero-padded index >= 1).
    import re as _re
    m = _re.match(r'^([a-zA-Z0-9._-]+)-(\d{2})$', name)
    if not m:
        return False
    safe_domain = m.group(1)
    # Reverse the wildcard.<x> -> *.<x> mapping that setup uses on disk.
    fqdn_candidates = [safe_domain]
    if safe_domain.startswith('wildcard.'):
        fqdn_candidates.append('*.' + safe_domain[len('wildcard.'):])
    try:
        with mysql_conn.cursor() as cur:
            placeholders = ','.join(['%s'] * len(fqdn_candidates))
            cur.execute(
                f"SELECT 1 FROM whp.site_domains WHERE fqdn IN ({placeholders}) LIMIT 1",
                fqdn_candidates
            )
            if cur.fetchone():
                return True
            cur.execute(
                f"SELECT 1 FROM whp.domains WHERE domain_name IN ({placeholders}) LIMIT 1",
                fqdn_candidates
            )
            return cur.fetchone() is not None
    except Exception:
        # Fail-safe: if the lookup fails, decline. The static allowlist
        # path will still match infrastructure containers (mysql, etc.).
        return False


def is_in_cooldown(sqlite_conn, target: str, cooldown_minutes: int) -> bool:
    """True if `target` (service name or container name) was remediated
    within the last cooldown_minutes. Read from the SQLite remediation_log
    so both engines share state."""
    try:
        cutoff = int(time.time()) - cooldown_minutes * 60
        row = sqlite_conn.execute(
            "SELECT 1 FROM remediation_log WHERE target = ? AND timestamp > ? LIMIT 1",
            (target, cutoff)
        ).fetchone()
        return row is not None
    except sqlite3.Error:
        return False


def attempt_remediation(*, sqlite_conn, mysql_conn,
                        category: str, signature: str,
                        ai_target: str = '', ai_confidence: str = '',
                        cooldown_minutes: int = 60,
                        require_high_confidence: bool = True,
                        dry_run: bool = True,
                        backup_load_multiplier: float = 4.0,
                        backup_swap_threshold: float = 90.0,
                        allowed_containers=None,
                        allowed_services=None) -> dict:
    """Decide whether to remediate, then either log the dry-run note or
    invoke the actual restart. Returns a dict with:
      action_taken: bool
      reason: short string explaining the decision
      command: the command line that ran (or would have run in dry-run)
      target: the resolved target name
    """
    allowed_containers = set(allowed_containers or [])
    allowed_services = set(allowed_services or [])
    result = {'action_taken': False, 'reason': '', 'command': '', 'target': ''}

    action = category_action(category)
    if action[0] == 'none':
        result['reason'] = f'category {category!r} is not auto-remediable'
        return result

    # Resolve target.
    kind = action[0]
    target = action[1]
    if target == 'auto':
        target = ai_target.strip()
        if not target:
            result['reason'] = 'AI did not specify a target'
            return result
    result['target'] = target

    # Confidence gate: only act when AI is high-confidence.
    if require_high_confidence and (ai_confidence or '').lower() != 'high':
        result['reason'] = f'confidence={ai_confidence or "(none)"} below required high'
        return result

    # Allowlist check.
    if kind == 'container' and target not in allowed_containers \
            and not _is_whp_customer_container(mysql_conn, target):
        result['reason'] = f"container {target!r} not in allowlist and not a WHP customer site"
        return result
    if kind == 'service' and target not in allowed_services:
        result['reason'] = f"service {target!r} not in allowlist"
        return result

    # Cooldown.
    if is_in_cooldown(sqlite_conn, target, cooldown_minutes):
        result['reason'] = f'{target} is in cooldown ({cooldown_minutes}m)'
        return result

    # Backup protection. The only way past this guard is system_at_risk(),
    # in which case we DO NOT restart anything -- we surface a special
    # action 'pause_backup' so the caller can pause the running backup
    # instead of bouncing services on top of it.
    if is_backup_running(mysql_conn):
        if system_at_risk(backup_load_multiplier, backup_swap_threshold):
            result['reason'] = ('backup running AND system_at_risk; '
                                'pausing backup instead of restart')
            result['command'] = 'PAUSE_BACKUP_REQUESTED'
            # Caller is responsible for actually pausing the queue worker;
            # we don't have a kill primitive here.
            return result
        result['reason'] = 'backup is currently running; refusing to restart anything'
        return result

    # Build the command.
    if kind == 'container':
        cmd = ['docker', 'restart', target]
    else:
        cmd = ['systemctl', 'restart', target]
    result['command'] = ' '.join(cmd)

    if dry_run:
        result['reason'] = f'[DRY RUN] would have run: {result["command"]}'
        _log_remediation(sqlite_conn, target=target, kind=kind,
                         command=result['command'], dry_run=True,
                         success=None, signature=signature)
        return result

    # Execute. Bounded subprocess; success measured by exit code 0.
    import subprocess
    try:
        proc = subprocess.run(cmd, capture_output=True, timeout=60, text=True)
        ok = (proc.returncode == 0)
    except Exception as e:
        ok = False
        proc = None
        result['reason'] = f'execution error: {e}'

    _log_remediation(sqlite_conn, target=target, kind=kind,
                     command=result['command'], dry_run=False,
                     success=ok, signature=signature,
                     stdout=(proc.stdout if proc else ''),
                     stderr=(proc.stderr if proc else result['reason']))
    if ok:
        result['action_taken'] = True
        result['reason'] = 'restart succeeded'
    else:
        result['reason'] = 'restart failed'
    return result


def _log_remediation(sqlite_conn, *, target, kind, command,
                     dry_run, success, signature,
                     stdout: str = '', stderr: str = ''):
    """Append a remediation_log row. Schema is created on demand here so
    poll-monitor and log-monitor share the same audit trail."""
    # SQLite's cursor.execute() only allows one statement per call; combining
    # CREATE TABLE + CREATE INDEX in one string raises "You can only execute
    # one statement at a time" and aborts the whole _log_remediation path.
    # Split into two calls (or use executescript). Two execs is clearer.
    sqlite_conn.execute("""
        CREATE TABLE IF NOT EXISTS remediation_log (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp INTEGER NOT NULL,
            target TEXT NOT NULL,
            action_type TEXT NOT NULL,
            command_run TEXT,
            stdout TEXT,
            stderr TEXT,
            exit_code INTEGER,
            success INTEGER,
            dry_run INTEGER DEFAULT 0,
            signature TEXT,
            error_id INTEGER
        )
    """)
    sqlite_conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_remediation_target_time "
        "ON remediation_log(target, timestamp)"
    )
    # Idempotent backfill for installs that pre-date a column. SQLite has
    # no IF NOT EXISTS for ADD COLUMN, so we PRAGMA-check and ALTER.
    cols = {row[1] for row in sqlite_conn.execute("PRAGMA table_info(remediation_log)").fetchall()}
    if 'dry_run' not in cols:
        try:
            sqlite_conn.execute("ALTER TABLE remediation_log ADD COLUMN dry_run INTEGER DEFAULT 0")
        except sqlite3.OperationalError:
            pass
    if 'signature' not in cols:
        try:
            sqlite_conn.execute("ALTER TABLE remediation_log ADD COLUMN signature TEXT")
        except sqlite3.OperationalError:
            pass
    sqlite_conn.execute(
        "INSERT INTO remediation_log "
        "(timestamp, target, action_type, command_run, stdout, stderr, "
        " success, dry_run, signature) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
        (int(time.time()), target, kind, command,
         (stdout or '')[:2000], (stderr or '')[:2000],
         (1 if success else 0) if success is not None else None,
         1 if dry_run else 0, signature)
    )
    sqlite_conn.commit()
