# NixOS Module for Honcho AI Memory Server
#
# This module provides systemd services for the Honcho API server and
# background deriver process, with configurable PostgreSQL and Redis backends.
#
# Usage in your NixOS configuration:
#
#   imports = [ inputs.m3ta-nixpkgs.nixosModules.default ];
#
#   m3ta.honcho = {
#     enable = true;
#     port = 8000;
#     host = "127.0.0.1";
#
#     # Database (PostgreSQL with pgvector)
#     database = {
#       connectionUri = "postgresql+psycopg://honcho:password@localhost:5432/honcho";
#     };
#
#     # Redis cache
#     cache = {
#       url = "redis://localhost:6379/0";
#     };
#
#     # LLM API keys (use agenix or sops-nix for secrets)
#     environmentFile = "/run/secrets/honcho-env";
#   };
#
# Using with m3ta.ports (recommended):
#
#   m3ta.ports = {
#     enable = true;
#     definitions = { honcho = 8000; };
#     currentHost = config.networking.hostName;
#   };
#
#   m3ta.honcho = {
#     enable = true;
#     port = config.m3ta.ports.get "honcho";
#   };
{
  config,
  lib,
  pkgs,
  ...
}:
with lib; let
  cfg = config.m3ta.honcho;

  # Python environment with honcho and FastAPI CLI
  pythonEnv = pkgs.python3.withPackages (ps:
    with ps; [
      cfg.package
    ]);

  # Start script for the API server
  startScript = pkgs.writeShellScript "honcho-start" ''
    set -e

    # Load environment file if specified
    ${optionalString (cfg.environmentFile != null) ''
      if [ -f "${cfg.environmentFile}" ]; then
        set -a
        source "${cfg.environmentFile}"
        set +a
      fi
    ''}

    # Create state directory
    mkdir -p ${cfg.stateDir}

    # Run database migrations
    ${pythonEnv}/bin/python -c "
import sys; sys.path.insert(0, '${cfg.package}/${pythonEnv.sitePackages}')
from scripts.provision_db import *
" 2>/dev/null || echo "Skipping database provisioning (script may not be available)"

    # Run the API server
    exec ${pythonEnv}/bin/fastapi run \
      --host ${cfg.host} \
      --port ${toString cfg.port} \
      ${cfg.package}/${pythonEnv.sitePackages}/src/main.py
  '';

  # Start script for the deriver background worker
  deriverScript = pkgs.writeShellScript "honcho-deriver" ''
    set -e

    # Load environment file if specified
    ${optionalString (cfg.environmentFile != null) ''
      if [ -f "${cfg.environmentFile}" ]; then
        set -a
        source "${cfg.environmentFile}"
        set +a
      fi
    ''}

    # Create state directory
    mkdir -p ${cfg.stateDir}

    # Run the deriver
    exec ${pythonEnv}/bin/python -m src.deriver
  '';
in {
  options.m3ta.honcho = {
    enable = mkEnableOption "Honcho AI memory server";

    package = mkOption {
      type = types.package;
      default = pkgs.honcho;
      defaultText = literalExpression "pkgs.honcho";
      description = "The honcho package to use.";
    };

    host = mkOption {
      type = types.str;
      default = "127.0.0.1";
      description = "Host address to bind the API server to.";
    };

    port = mkOption {
      type = types.port;
      default = 8000;
      description = "Port to run the API server on.";
    };

    stateDir = mkOption {
      type = types.path;
      default = "/var/lib/honcho";
      description = "Directory to store honcho data and state.";
    };

    user = mkOption {
      type = types.str;
      default = "honcho";
      description = "User account under which honcho runs.";
    };

    group = mkOption {
      type = types.str;
      default = "honcho";
      description = "Group under which honcho runs.";
    };

    environmentFile = mkOption {
      type = types.nullOr types.path;
      default = null;
      description = ''
        Environment file containing configuration and secrets.
        This file should contain KEY=value pairs, one per line.
        Use this for API keys, database credentials, and other secrets.
        See the honcho .env.template for available variables.
      '';
      example = "/run/secrets/honcho-env";
    };

    # Database Configuration
    database = {
      connectionUri = mkOption {
        type = types.str;
        default = "postgresql+psycopg://honcho:honcho@localhost:5432/honcho";
        description = ''
          PostgreSQL connection URI with pgvector support.
          Must use postgresql+psycopg prefix for SQLAlchemy compatibility.
        '';
        example = "postgresql+psycopg://honcho:password@localhost:5432/honcho";
      };

      schema = mkOption {
        type = types.str;
        default = "public";
        description = "Database schema to use.";
      };

      poolSize = mkOption {
        type = types.int;
        default = 10;
        description = "Connection pool size.";
      };

      maxOverflow = mkOption {
        type = types.int;
        default = 20;
        description = "Maximum connection pool overflow.";
      };
    };

    # Cache Configuration
    cache = {
      url = mkOption {
        type = types.str;
        default = "redis://localhost:6379/0";
        description = ''
          Redis cache URL.
          Set suppress=true to suppress connection errors: redis://localhost:6379/0?suppress=true
        '';
        example = "redis://localhost:6379/0?suppress=true";
      };
    };

    # Deriver (background worker) Configuration
    deriver = {
      enable = mkOption {
        type = types.bool;
        default = true;
        description = ''
          Enable the honcho deriver background worker.
          Processes asynchronous tasks like representation updates and session summarization.
        '';
      };

      workers = mkOption {
        type = types.int;
        default = 1;
        description = "Number of deriver worker processes.";
      };

      provider = mkOption {
        type = types.str;
        default = "google";
        description = "LLM provider for the deriver.";
      };

      model = mkOption {
        type = types.str;
        default = "gemini-2.5-flash-lite";
        description = "LLM model for the deriver.";
      };
    };

    # Logging
    logLevel = mkOption {
      type = types.enum ["CRITICAL" "ERROR" "WARNING" "INFO" "DEBUG"];
      default = "INFO";
      description = "Logging level for the server.";
    };

    # Authentication
    auth = {
      enable = mkOption {
        type = types.bool;
        default = false;
        description = "Enable JWT authentication.";
      };

      jwtSecretFile = mkOption {
        type = types.nullOr types.path;
        default = null;
        description = ''
          Path to file containing the JWT secret key.
          Required when auth is enabled.
        '';
        example = "/run/secrets/honcho-jwt-secret";
      };
    };
  };

  config = mkIf cfg.enable {
    # Create user and group
    users.users.${cfg.user} = {
      isSystemUser = true;
      group = cfg.group;
      description = "Honcho service user";
      home = cfg.stateDir;
      createHome = true;
    };

    users.groups.${cfg.group} = {};

    # API Server systemd service
    systemd.services.honcho = {
      description = "Honcho AI Memory API Server";
      after = ["network.target"];
      wantedBy = ["multi-user.target"];

      serviceConfig = {
        Type = "simple";
        User = cfg.user;
        Group = cfg.group;
        ExecStart = startScript;
        Restart = "on-failure";
        RestartSec = "5s";

        # Security hardening
        NoNewPrivileges = true;
        PrivateTmp = true;
        ProtectSystem = "strict";
        ProtectHome = true;
        ReadWritePaths = [cfg.stateDir];
        ProtectKernelTunables = true;
        ProtectKernelModules = true;
        ProtectControlGroups = true;
        RestrictRealtime = true;
        RestrictNamespaces = true;
        LockPersonality = true;
        MemoryDenyWriteExecute = false; # Python needs this
        RestrictAddressFamilies = ["AF_UNIX" "AF_INET" "AF_INET6"];
      };

      environment =
        {
          PYTHONUNBUFFERED = "1";
          DB_CONNECTION_URI = cfg.database.connectionUri;
          DB_SCHEMA = cfg.database.schema;
          DB_POOL_SIZE = toString cfg.database.poolSize;
          DB_MAX_OVERFLOW = toString cfg.database.maxOverflow;
          CACHE_URL = cfg.cache.url;
          LOG_LEVEL = cfg.logLevel;
          AUTH_USE_AUTH =
            if cfg.auth.enable
            then "true"
            else "false";
          DERIVER_ENABLED =
            if cfg.deriver.enable
            then "true"
            else "false";
        }
        // optionalAttrs (cfg.auth.jwtSecretFile != null) {
          AUTH_JWT_SECRET_FILE = cfg.auth.jwtSecretFile;
        };
    };

    # Deriver background worker systemd service
    systemd.services.honcho-deriver = mkIf cfg.deriver.enable {
      description = "Honcho Deriver Background Worker";
      after = ["network.target"];
      wantedBy = ["multi-user.target"];

      serviceConfig = {
        Type = "simple";
        User = cfg.user;
        Group = cfg.group;
        ExecStart = deriverScript;
        Restart = "on-failure";
        RestartSec = "5s";

        # Security hardening
        NoNewPrivileges = true;
        PrivateTmp = true;
        ProtectSystem = "strict";
        ProtectHome = true;
        ReadWritePaths = [cfg.stateDir];
        ProtectKernelTunables = true;
        ProtectKernelModules = true;
        ProtectControlGroups = true;
        RestrictRealtime = true;
        RestrictNamespaces = true;
        LockPersonality = true;
        MemoryDenyWriteExecute = false; # Python needs this
        RestrictAddressFamilies = ["AF_UNIX" "AF_INET" "AF_INET6"];
      };

      environment =
        {
          PYTHONUNBUFFERED = "1";
          DB_CONNECTION_URI = cfg.database.connectionUri;
          DB_SCHEMA = cfg.database.schema;
          CACHE_URL = cfg.cache.url;
          LOG_LEVEL = cfg.logLevel;
          DERIVER_ENABLED = "true";
          DERIVER_WORKERS = toString cfg.deriver.workers;
          DERIVER_PROVIDER = cfg.deriver.provider;
          DERIVER_MODEL = cfg.deriver.model;
          METRICS_ENABLED = "true";
        }
        // optionalAttrs (cfg.auth.jwtSecretFile != null) {
          AUTH_JWT_SECRET_FILE = cfg.auth.jwtSecretFile;
        };
    };

    # Open firewall port if binding to non-localhost
    networking.firewall.allowedTCPPorts = mkIf (cfg.host != "127.0.0.1" && cfg.host != "localhost") [cfg.port];
  };
}