import os
import sys
import json
import socket
import struct
import argparse
import subprocess
import warnings
from datetime import datetime, timezone, timedelta
from typing import Optional, Dict, Any, List

# Suppress cryptography / pgpy deprecation notices
warnings.filterwarnings("ignore")

import pgpy

CONFIG_FILE_NAME = "client_config.json"
STATE_FILE_NAME = "client_state.json"


def get_state_path(config_path: str, custom_state_path: Optional[str] = None) -> str:
    if custom_state_path:
        return custom_state_path
    config_dir = os.path.dirname(os.path.abspath(config_path))
    return os.path.join(config_dir, STATE_FILE_NAME)


def load_state(state_path: str) -> dict:
    if os.path.exists(state_path):
        try:
            with open(state_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"[!] Warning: Failed to read state file '{state_path}': {e}")
            return {}
    return {}


def save_state(state_path: str, state: dict):
    try:
        temp_path = f"{state_path}.tmp"
        with open(temp_path, "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2)
        os.replace(temp_path, state_path)
    except Exception as e:
        print(f"[!] Warning: Could not save client state to '{state_path}': {e}")


def commit_state(state: dict, state_path: str):
    if "new_last_cursor" in state:
        val = state.pop("new_last_cursor")
        if val:
            state["last_cursor"] = val
    if "new_last_timestamp_us" in state:
        val = state.pop("new_last_timestamp_us")
        if val:
            state["last_timestamp_us"] = val
    if "new_sent_cursors" in state:
        state["sent_cursors"] = state.pop("new_sent_cursors")
    save_state(state_path, state)


def load_config(config_path: str = CONFIG_FILE_NAME):
    if not os.path.exists(config_path):
        raise FileNotFoundError(
            f"Client configuration file not found at: {config_path}\n"
            f"Generate one from the server using: python Server.py --create-client-config --client-out {config_path}"
        )
    with open(config_path, "r", encoding="utf-8") as f:
        return json.load(f)


def get_machine_identifier() -> str:
    """
    Returns the hostname of the machine sending the logs,
    and appends the network/DNS domain if available.
    """
    # 1. Try fully-qualified domain name (FQDN)
    fqdn = socket.getfqdn()
    if fqdn and "." in fqdn and not fqdn.startswith("localhost"):
        return fqdn

    hostname = socket.gethostname()

    # 2. Check /etc/resolv.conf domain or search directive
    try:
        if os.path.exists("/etc/resolv.conf"):
            with open("/etc/resolv.conf", "r", encoding="utf-8") as f:
                for line in f:
                    parts = line.strip().split()
                    if parts and parts[0] in ["domain", "search"] and len(parts) > 1:
                        domain = parts[1]
                        if domain and not domain.startswith("."):
                            return f"{hostname}.{domain}"
    except Exception:
        pass

    # 3. Try reverse DNS lookup
    try:
        host_ip = socket.gethostbyname(hostname)
        canonical_name = socket.gethostbyaddr(host_ip)[0]
        if canonical_name and "." in canonical_name and not canonical_name.startswith("localhost"):
            return canonical_name
    except Exception:
        pass

    return hostname


def get_recent_linux_logs(hours: int = 24, state: Optional[dict] = None) -> list:
    """
    Collects info, warnings, and errors from systemd journalctl over the lookback window.
    Edge Filtering: Retains INFO, WARNING, and ERROR. Strips DEBUG (priority 7) and skips events older than lookback window (default: 24h).
    State Tracking: Skips events older than lookback window (default 24h) and events
    that have already been sent in previous runs.
    """
    last_cursor = None
    last_timestamp_us = 0.0
    sent_cursors = set()
    if state:
        last_cursor = state.get("last_cursor")
        try:
            last_timestamp_us = float(state.get("last_timestamp_us", 0))
        except (ValueError, TypeError):
            last_timestamp_us = 0.0
        sent_cursors = set(state.get("sent_cursors", []))

    cmd = ["journalctl", "--since", f"{hours} hours ago", "-p", "info", "--output=json"]
    result = None

    if last_cursor:
        cmd_with_cursor = ["journalctl", "--since", f"{hours} hours ago", "--after-cursor", str(last_cursor), "-p", "info", "--output=json"]
        try:
            res = subprocess.run(cmd_with_cursor, capture_output=True, text=True, check=False)
            if res.returncode == 0:
                result = res
        except FileNotFoundError:
            print("[!] journalctl command not found. Ensure this script runs on a systemd-enabled Linux system.")
            return []
        except Exception:
            pass

    if result is None:
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=False)
        except FileNotFoundError:
            print("[!] journalctl command not found. Ensure this script runs on a systemd-enabled Linux system.")
            return []
        except Exception as e:
            print(f"[!] Error running journalctl: {e}")
            return []

    logs = []
    machine_id = get_machine_identifier()
    cutoff_epoch_us = (datetime.now(timezone.utc) - timedelta(hours=hours)).timestamp() * 1_000_000

    newest_cursor = None
    newest_timestamp_us = last_timestamp_us
    collected_cursors = []

    for line in result.stdout.splitlines():
        line_str = line.strip()
        if not line_str:
            continue
        try:
            entry = json.loads(line_str)
            entry_cursor = entry.get("__CURSOR")
            entry_ts_us_raw = entry.get("__REALTIME_TIMESTAMP")

            entry_ts_us = 0.0
            if entry_ts_us_raw:
                try:
                    entry_ts_us = float(entry_ts_us_raw)
                except (ValueError, TypeError):
                    pass

            # 1. Skip entries older than lookback window (default: 24h)
            if entry_ts_us and entry_ts_us < cutoff_epoch_us:
                continue

            # 2. Skip already sent events
            if entry_cursor and (entry_cursor in sent_cursors or entry_cursor == last_cursor):
                continue
            if last_timestamp_us > 0 and entry_ts_us > 0 and entry_ts_us < last_timestamp_us:
                continue

            # Advance newest tracking for new entries
            if entry_cursor:
                newest_cursor = entry_cursor
                collected_cursors.append(entry_cursor)
            if entry_ts_us > newest_timestamp_us:
                newest_timestamp_us = entry_ts_us

            priority = int(entry.get("PRIORITY", "6"))
            # Priority 0: Emerg, 1: Alert, 2: Crit, 3: Err (-> ERROR)
            # Priority 4: Warning, 5: Notice (-> WARNING)
            # Priority 6: Info (-> INFO)
            # Priority 7: Debug (skip)
            if priority > 6:
                continue

            if priority <= 3:
                sev = "ERROR"
            elif priority in (4, 5):
                sev = "WARNING"
            else:
                sev = "INFO"

            logs.append({
                "server": machine_id,
                "os_type": "linux",
                "signature": entry.get("SYSLOG_IDENTIFIER", "unknown"),
                "severity": sev,
                "message": entry.get("MESSAGE", "")[:2048]
            })
        except (json.JSONDecodeError, ValueError):
            continue

    if state is not None:
        state["new_last_cursor"] = newest_cursor or last_cursor
        state["new_last_timestamp_us"] = max(newest_timestamp_us, last_timestamp_us)
        state["new_sent_cursors"] = (list(sent_cursors) + collected_cursors)[-1000:]
        state["last_run_timestamp"] = datetime.now(timezone.utc).isoformat()

    return logs


def send_encrypted_logs_over_socket(config: dict, logs: list):
    """
    Encrypts the payload using the server's OpenPGP public key and streams
    over an authenticated TCP socket.
    """
    server_host = config["server_host"]
    server_port = int(config["server_port"])
    auth_token = config["auth_token"]
    pub_key_armored = config["server_public_key"]
    expected_fp = config.get("server_fingerprint", "").replace(" ", "").upper()

    # Load and verify server public key
    pub_key, _ = pgpy.PGPKey.from_blob(pub_key_armored)
    actual_fp = str(pub_key.fingerprint).replace(" ", "").upper()
    if expected_fp and actual_fp != expected_fp:
        raise ValueError(
            f"Server fingerprint mismatch! Expected {expected_fp}, but key has {actual_fp}."
        )

    machine_id = get_machine_identifier()

    # Prepare batch
    payload = {
        "server": machine_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "logs": logs
    }
    payload_json = json.dumps(payload)

    # Encrypt payload with server's encryption-only key
    pgp_msg = pgpy.PGPMessage.new(payload_json)
    encrypted_msg = pub_key.encrypt(pgp_msg)
    encrypted_armored = str(encrypted_msg)

    # Envelope with socket authentication header
    envelope = {
        "auth_token": auth_token,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "encrypted_payload": encrypted_armored
    }
    envelope_bytes = json.dumps(envelope).encode("utf-8")

    # Connect over TCP socket and transmit with 4-byte length prefix framing
    print(f"[*] Connecting to LOGAR server at {server_host}:{server_port} over secure TCP socket...")
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.settimeout(15.0)
        sock.connect((server_host, server_port))

        # Send frame: length (4 bytes big-endian) + envelope
        frame = struct.pack(">I", len(envelope_bytes)) + envelope_bytes
        sock.sendall(frame)

        # Receive response length
        resp_len_bytes = sock.recv(4)
        if not resp_len_bytes:
            raise ConnectionError("Server closed connection without response.")

        resp_len = struct.unpack(">I", resp_len_bytes)[0]
        resp_bytes = bytearray()
        while len(resp_bytes) < resp_len:
            chunk = sock.recv(min(4096, resp_len - len(resp_bytes)))
            if not chunk:
                break
            resp_bytes.extend(chunk)

        response = json.loads(resp_bytes.decode("utf-8"))
        print(f"[+] Server response: {response}")
        return response


def main():
    parser = argparse.ArgumentParser(description="LOGAR Linux Edge Log Forwarder with State Tracking")
    parser.add_argument("--config", default=CONFIG_FILE_NAME, help="Path to client_config.json")
    parser.add_argument("--hours", type=int, default=24, help="Lookback window in hours for journalctl logs (default: 24)")
    parser.add_argument("--state-file", default=None, help="Path to state tracking file (default: client_state.json next to config)")
    parser.add_argument("--no-state", action="store_true", help="Disable state tracking and send all events matching lookback window")
    args = parser.parse_args()

    try:
        config = load_config(args.config)
    except Exception as e:
        print(f"[!] Configuration error: {e}")
        sys.exit(1)

    state_path = get_state_path(args.config, args.state_file)
    state = None if args.no_state else load_state(state_path)

    machine_id = get_machine_identifier()
    print(f"[*] Edge Forwarder Node: {machine_id}")
    if state and ("last_cursor" in state or "last_timestamp_us" in state):
        print(f"[*] State tracking active: resuming after previous cursor/timestamp (state file: {state_path})")
    elif not args.no_state:
        print(f"[*] State tracking initialized (state file: {state_path})")

    print(f"[*] Scanning Linux journalctl for unsent entries (last {args.hours} hours)...")
    candidate_logs = get_recent_linux_logs(hours=args.hours, state=state)
    print(f"[*] Found {len(candidate_logs)} unsent candidate entries (INFO to ERROR, entries > {args.hours}h and already-sent skipped).")

    if not candidate_logs:
        print("[*] No new unsent events to transmit.")
        if state is not None:
            commit_state(state, state_path)
        return

    try:
        resp = send_encrypted_logs_over_socket(config, candidate_logs)
        if state is not None and resp and resp.get("status") == "success":
            commit_state(state, state_path)
            print(f"[+] State successfully committed to {state_path}")
    except Exception as e:
        print(f"[!] Failed to stream logs to server: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()