import os
import re
import json
import ipaddress
import subprocess
import requests
import smtplib
from email.message import EmailMessage
from datetime import date, timedelta
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()
]

# --- Multi-day repeat-offender tracking (persisted on the /data mount) ---
STATE_FILE = os.getenv("STATE_FILE", "/data/offender_state.json")
# An IP is called out as a "repeat offender" once it's shown up on this many distinct days
REPEAT_OFFENDER_THRESHOLD_DAYS = int(os.getenv("REPEAT_OFFENDER_THRESHOLD_DAYS", 3))
# IPs not seen again after this many days are dropped from state, so the file doesn't grow forever
OFFENDER_RETENTION_DAYS = int(os.getenv("OFFENDER_RETENTION_DAYS", 30))

# --- 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.:a-fA-F]+)\]:\d+: (\d{3} \S+ [^;]+)')
DNSBL_RE = re.compile(r'dnsblog\[\d+\]: addr')
RELAY_DENIED_RE = re.compile(r'reject: RCPT from \S+\[([\d.:a-fA-F]+)\]: 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()
    postscreen_reject_ips = Counter()
    dnsbl_hits = 0
    relay_denied = []
    relay_denied_ips = Counter()
    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_reject_ips[m.group(1)] += 1
            postscreen_rejects[m.group(2)] += 1
        if DNSBL_RE.search(line):
            dnsbl_hits += 1
        if m := RELAY_DENIED_RE.search(line):
            relay_denied_ips[m.group(1)] += 1
            relay_denied.append(m.group(2))
        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

    # Union of IPs worth remembering across days: anything that got rejected,
    # probed with junk protocol data, or tried to abuse the relay. Routine
    # DNSBL-only connects are deliberately excluded -- too noisy/low-signal
    # to track individually (see prior discussion on scanner noise).
    suspicious_ips_today = Counter()
    for ip, count in nonsmtp_by_ip.items():
        suspicious_ips_today[ip] += count
    for ip, count in postscreen_reject_ips.items():
        suspicious_ips_today[ip] += count
    for ip, count in relay_denied_ips.items():
        suspicious_ips_today[ip] += count
    for user, method, ip in untrusted_logins:
        suspicious_ips_today[ip] += 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, suspicious_ips_today


def load_offender_state():
    if not os.path.exists(STATE_FILE):
        debug_print(f"No existing state file at {STATE_FILE}, starting fresh.")
        return {}
    try:
        with open(STATE_FILE, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError) as e:
        debug_print(f"Could not read state file ({e}), starting fresh.")
        return {}


def save_offender_state(state):
    try:
        os.makedirs(os.path.dirname(STATE_FILE), exist_ok=True)
        with open(STATE_FILE, "w") as f:
            json.dump(state, f, indent=2, sort_keys=True)
    except OSError as e:
        debug_print(f"Could not write state file: {e}")


def update_offender_state(state, suspicious_ips_today, report_date_str):
    """
    Records which IPs were suspicious today, merges into the persisted
    history, prunes IPs not seen recently, and returns (new_state, repeat_offenders).
    repeat_offenders is a list of dicts sorted by distinct-day count descending.
    """
    for ip, count in suspicious_ips_today.items():
        entry = state.setdefault(ip, {"days": [], "hit_counts": {}})
        if report_date_str not in entry["days"]:
            entry["days"].append(report_date_str)
        entry["hit_counts"][report_date_str] = entry["hit_counts"].get(report_date_str, 0) + count
        entry["last_seen"] = report_date_str

    # Prune anything not seen within the retention window
    cutoff = date.fromisoformat(report_date_str) - timedelta(days=OFFENDER_RETENTION_DAYS)
    pruned_state = {}
    for ip, entry in state.items():
        last_seen = date.fromisoformat(entry.get("last_seen", report_date_str))
        if last_seen >= cutoff:
            pruned_state[ip] = entry

    repeat_offenders = []
    for ip, entry in pruned_state.items():
        distinct_days = len(entry["days"])
        if distinct_days >= REPEAT_OFFENDER_THRESHOLD_DAYS:
            repeat_offenders.append({
                "ip": ip,
                "distinct_days": distinct_days,
                "first_seen": min(entry["days"]),
                "last_seen": entry["last_seen"],
                "total_hits": sum(entry["hit_counts"].values()),
            })
    repeat_offenders.sort(key=lambda x: x["distinct_days"], reverse=True)

    return pruned_state, repeat_offenders


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, suspicious_ips_today = analyze_logs(logs)

    report_date_str = (date.today() - timedelta(days=1)).isoformat()
    state = load_offender_state()
    state, repeat_offenders = update_offender_state(state, suspicious_ips_today, report_date_str)
    save_offender_state(state)

    if repeat_offenders:
        needs_attention = True
        lines = [f"Repeat offenders (seen probing/rejected on {REPEAT_OFFENDER_THRESHOLD_DAYS}+ distinct days -- consider a firewall block):"]
        for o in repeat_offenders:
            lines.append(
                f"  - {o['ip']}: seen on {o['distinct_days']} distinct days "
                f"(first {o['first_seen']}, last {o['last_seen']}), {o['total_hits']} total hits"
            )
        facts_text = facts_text + "\n" + "\n".join(lines)
    else:
        facts_text = facts_text + f"\nRepeat offenders (seen on {REPEAT_OFFENDER_THRESHOLD_DAYS}+ distinct days): 0"

    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)