import os
import re
import ipaddress
import subprocess
import requests
import smtplib
from email.message import EmailMessage
import time
import schedule
from collections import Counter, defaultdict

# --- CONFIGURATION (Loaded from environment variables) ---
OLLAMA_API = os.getenv("OLLAMA_API_URL", "http://localhost:11434/api/generate")
MODEL = os.getenv("OLLAMA_MODEL", "llama3.2:1b")
OLLAMA_TIMEOUT = int(os.getenv("OLLAMA_TIMEOUT_SECONDS", 300))

AI_SYSTEM_PROMPT = os.getenv("AI_SYSTEM_PROMPT", (
    "You are a Linux sysadmin writing a brief daily report for a technically "
    "experienced operator. You are given PRE-EXTRACTED, ALREADY-VERIFIED FACTS "
    "below -- do not invent, infer, or add any event, IP, or attack that is not "
    "explicitly listed. If a category shows 0 or is absent, say nothing happened "
    "in that category. Do not editorialize or add generic security recommendations "
    "unless a fact explicitly calls for one."
))

MY_EMAIL = os.getenv("REPORT_TO_EMAIL")
FROM_EMAIL = os.getenv("REPORT_FROM_EMAIL")
SMTP_SERVER = os.getenv("SMTP_SERVER", "127.0.0.1")
SMTP_PORT = int(os.getenv("SMTP_PORT", 25))
SMTP_USER = os.getenv("SMTP_USER", "")
SMTP_PASS = os.getenv("SMTP_PASS", "")

JOURNAL_UNIT = os.getenv("JOURNAL_UNIT", "mailserver")

# Networks/IPs that are allowed to log in -- anything outside this is flagged
TRUSTED_LOGIN_NETWORKS = [
    n.strip() for n in os.getenv("TRUSTED_LOGIN_NETWORKS", "10.0.0.0/24").split(",") if n.strip()
]

# --- EXECUTION TOGGLES ---
DEBUG = os.getenv("DEBUG", "false").lower() in ("true", "1", "yes")
RUN_NOW = os.getenv("RUN_NOW", "false").lower() in ("true", "1", "yes")
REPORT_TIME = os.getenv("REPORT_TIME", "08:00")


def debug_print(message):
    """Prints only if DEBUG=true in .env"""
    if DEBUG:
        print(f"[DEBUG] {message}")


def get_yesterdays_logs():
    print(f"Fetching logs for '{JOURNAL_UNIT}' from yesterday...")
    debug_print(f"Running command: journalctl -u {JOURNAL_UNIT} --since yesterday --until today")

    result = subprocess.run(
        ['journalctl', '-u', JOURNAL_UNIT, '--since', 'yesterday', '--until', 'today'],
        capture_output=True, text=True
    )
    lines = result.stdout.splitlines()
    debug_print(f"Fetched {len(lines)} total lines from journalctl.")
    return lines


def _is_trusted(ip):
    try:
        addr = ipaddress.ip_address(ip)
    except ValueError:
        return False
    for net in TRUSTED_LOGIN_NETWORKS:
        try:
            if addr in ipaddress.ip_network(net, strict=False):
                return True
        except ValueError:
            continue
    return False


# --- Regexes for structured extraction (mail server specific) ---
CONNECT_RE = re.compile(r'postscreen\[\d+\]: CONNECT from \[([\d.:a-fA-F]+)\]')
REJECT_RE = re.compile(r'postscreen\[\d+\]: NOQUEUE: reject: RCPT from \[[^\]]+\]:\d+: (\d{3} \S+ [^;]+)')
DNSBL_RE = re.compile(r'dnsblog\[\d+\]: addr')
RELAY_DENIED_RE = re.compile(r'reject: RCPT from \S+: 554 [\d.]+ <[^>]+>: Relay access denied.*from=<([^>]+)>')
LOGIN_RE = re.compile(r'imap-login: Login: user=<([^>]+)>, method=(\S+), rip=([\d.:a-fA-F]+)')
AUTH_FAIL_RE = re.compile(r'auth failed|authentication failure|invalid password', re.IGNORECASE)
SENT_RE = re.compile(r'status=sent')
NONSMTP_IP_RE = re.compile(r'non-SMTP command from (?:unknown\[([\d.]+)\]|[\w.\-]+\[([\d.]+)\])')
SENDER_AUTH_ANOMALY_RE = re.compile(r'opendkim.*key retrieval failed|opendkim.*failed to parse|policyd-spf.*Temperror|policyd-spf.*Fail')
GENERIC_WARNING_RE = re.compile(r'warning: (.*)')
HOSTNAME_NORESOLVE_RE = re.compile(r'hostname .* does not resolve')


def analyze_logs(lines):
    """
    Deterministically extract structured, verifiable facts from raw mail server
    log lines. This replaces naive keyword grepping -- the LLM only narrates
    these facts, it never has to "reason" over raw log text, which is what a
    small local model reliably fails at.
    """
    total_lines = len(lines)

    postscreen_rejects = Counter()
    dnsbl_hits = 0
    relay_denied = []
    logins = defaultdict(list)
    untrusted_logins = []
    auth_failures = []
    mail_delivered = 0
    nonsmtp_by_ip = Counter()
    sender_auth_anomalies = 0
    other_warnings = Counter()
    hostname_noresolve_count = 0

    for line in lines:
        if m := REJECT_RE.search(line):
            postscreen_rejects[m.group(1)] += 1
        if DNSBL_RE.search(line):
            dnsbl_hits += 1
        if m := RELAY_DENIED_RE.search(line):
            relay_denied.append(m.group(1))
        if m := LOGIN_RE.search(line):
            user, method, ip = m.group(1), m.group(2), m.group(3)
            logins[user].append(ip)
            if not _is_trusted(ip):
                untrusted_logins.append((user, method, ip))
        if AUTH_FAIL_RE.search(line):
            auth_failures.append(line.strip())
        if SENT_RE.search(line):
            mail_delivered += 1
        if m := NONSMTP_IP_RE.search(line):
            ip = m.group(1) or m.group(2)
            nonsmtp_by_ip[ip] += 1
        elif HOSTNAME_NORESOLVE_RE.search(line):
            hostname_noresolve_count += 1
        elif SENDER_AUTH_ANOMALY_RE.search(line):
            sender_auth_anomalies += 1
        elif m := GENERIC_WARNING_RE.search(line):
            other_warnings[m.group(1).strip()] += 1

    facts = []
    facts.append(f"Total log lines: {total_lines}")
    facts.append(f"Postscreen automatic rejections (bots/scanners blocked before reaching mailbox): {sum(postscreen_rejects.values())}")
    for reason, count in postscreen_rejects.most_common():
        facts.append(f"  - {count}x: {reason}")
    facts.append(f"DNSBL blacklist hits recorded (routine scoring of inbound connections): {dnsbl_hits}")

    facts.append(f"Open-relay abuse attempts blocked (relay access denied): {len(relay_denied)}")
    for env in relay_denied:
        facts.append(f"  - envelope-from: {env}")

    total_logins = sum(len(v) for v in logins.values())
    facts.append(f"Successful mailbox logins: {total_logins}")
    for user, ips in logins.items():
        facts.append(f"  - user={user}, {len(ips)} session(s), source IP(s): {sorted(set(ips))}")

    facts.append(f"Logins from UNTRUSTED/unexpected source IPs (outside {TRUSTED_LOGIN_NETWORKS}): {len(untrusted_logins)}")
    for user, method, ip in untrusted_logins:
        facts.append(f"  - user={user}, method={method}, source IP={ip}  <-- FLAG THIS, review immediately")

    facts.append(f"Authentication FAILURES (should normally be 0; any value above 0 is a possible brute-force indicator): {len(auth_failures)}")
    for l in auth_failures[:15]:
        facts.append(f"  - {l}")

    facts.append(f"Legitimate messages delivered to mailbox: {mail_delivered}")

    if nonsmtp_by_ip:
        facts.append("Junk/non-SMTP protocol probes on port 25, by source IP (typical internet-wide scanner noise, e.g. raw TLS/HTTP bytes sent to the SMTP port -- NOT a login attempt):")
        for ip, count in nonsmtp_by_ip.most_common():
            facts.append(f"  - {ip}: {count} probe(s)")

    facts.append(f"Inbound sender DKIM/SPF anomalies (misconfiguration on the SENDING domain's side, not this server): {sender_auth_anomalies}")
    facts.append(f"Reverse-DNS 'hostname does not resolve' warnings (routine, from senders with no PTR record): {hostname_noresolve_count}")

    if other_warnings:
        facts.append(f"Other warning lines not otherwise categorized ({sum(other_warnings.values())} total):")
        for w, count in other_warnings.most_common(10):
            facts.append(f"  - ({count}x) {w}")

    needs_attention = bool(auth_failures) or bool(untrusted_logins) or bool(relay_denied) or bool(other_warnings)

    return total_lines, "\n".join(facts), needs_attention


def get_ai_summary(facts_text):
    print("Asking AI to write a summary from pre-extracted facts...")

    prompt = (
        f"{AI_SYSTEM_PROMPT}\n\n"
        f"FACTS:\n{facts_text}\n\n"
        f"Write a brief report (max 150 words): one line on overall status, one "
        f"line on mail delivered, then list any items flagged above only if their "
        f"count is greater than zero. Do not repeat the raw fact list verbatim; "
        f"synthesize it into prose."
    )

    payload = {
        "model": MODEL,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.1,
            "num_predict": 400,
        },
    }

    debug_print(f"Sending prompt to Ollama at {OLLAMA_API} using model '{MODEL}'")
    debug_print(f"Total prompt length: {len(prompt)} characters.")

    try:
        response = requests.post(OLLAMA_API, json=payload, timeout=OLLAMA_TIMEOUT)
        response.raise_for_status()
        debug_print(f"AI response status code: {response.status_code}")
        data = response.json()
        return data.get('response', 'Error: AI returned empty response.').strip()
    except requests.exceptions.Timeout:
        debug_print(f"AI request timed out after {OLLAMA_TIMEOUT}s")
        return f"AI summary unavailable: Ollama did not respond within {OLLAMA_TIMEOUT}s. See raw facts below."
    except Exception as e:
        debug_print(f"AI Connection Exception: {str(e)}")
        return f"AI summary unavailable ({str(e)}). See raw facts below."


def send_email_report(total_traffic, ai_summary, facts_text, needs_attention):
    if not MY_EMAIL or not FROM_EMAIL:
        print("Skipping email: REPORT_TO_EMAIL or REPORT_FROM_EMAIL not set.")
        return

    print(f"Sending email report to {MY_EMAIL}...")
    debug_print(f"SMTP Server: {SMTP_SERVER}:{SMTP_PORT}")
    debug_print(f"SMTP Auth User: {'<configured>' if SMTP_USER else '<none>'}")

    tag = "[ATTENTION] " if needs_attention else "[OK] "

    msg = EmailMessage()
    msg['Subject'] = f"{tag}Daily AI Summary: {JOURNAL_UNIT}"
    msg['From'] = FROM_EMAIL
    msg['To'] = MY_EMAIL

    body = (
        f"DAILY SERVER REPORT: {JOURNAL_UNIT}\n"
        f"----------------------------------------\n"
        f"Total Log Lines Processed: {total_traffic}\n\n"
        f"AI Summary:\n{ai_summary}\n\n"
        f"----------------------------------------\n"
        f"Raw extracted facts (ground truth -- the AI summary above is generated "
        f"from this and should match it):\n{facts_text}\n"
    )

    msg.set_content(body)

    try:
        server = smtplib.SMTP(SMTP_SERVER, SMTP_PORT)
        if SMTP_USER:
            debug_print("Attempting STARTTLS and Login...")
            server.starttls()
            server.login(SMTP_USER, SMTP_PASS)

        debug_print("Transmitting email payload...")
        server.send_message(msg)
        server.quit()
        print("Email sent successfully!")
    except Exception as e:
        print(f"Failed to send email: {e}")
        debug_print(f"Email Exception: {str(e)}")


def job():
    print(f"\n--- Starting log analysis for {JOURNAL_UNIT} ---")
    logs = get_yesterdays_logs()
    total_traffic, facts_text, needs_attention = analyze_logs(logs)
    debug_print(f"Extracted facts:\n{facts_text}")
    summary = get_ai_summary(facts_text)
    send_email_report(total_traffic, summary, facts_text, needs_attention)
    print("--- Analysis complete. ---")


if __name__ == "__main__":
    print("Log Analyzer Container starting up...")

    if DEBUG:
        print("[DEBUG] Debug mode is ENABLED.")

    # 1. Register the permanent daily schedule first
    schedule.every().day.at(REPORT_TIME).do(job)
    print(f"Job permanently scheduled to run daily at {REPORT_TIME}.")

    # 2. Handle the immediate run if requested
    if RUN_NOW:
        print("RUN_NOW is enabled. Executing immediate startup run...")
        job()
        print("--- Immediate run finished. Returning to standard schedule. ---")
    else:
        print("RUN_NOW is disabled. Skipping immediate run.")

    # 3. Enter the permanent background loop
    print(f"Entering background scheduler loop. Waiting for {REPORT_TIME}...")
    while True:
        schedule.run_pending()
        time.sleep(60)