# NixOS Module for Mem0 REST API Server
#
# This module provides a systemd service for the Mem0 REST API server,
# allowing you to run mem0 as a system service with configurable vector storage.
#
# Usage in your NixOS configuration:
#
#   # In your flake.nix or configuration.nix:
#   imports = [ inputs.m3ta-nixpkgs.nixosModules.default ];
#
#   m3ta.mem0 = {
#     enable = true;
#     port = 8000;
#     host = "127.0.0.1";
#
#     # LLM Configuration
#     llm = {
#       provider = "openai";
#       apiKeyFile = "/run/secrets/openai-api-key";  # Use agenix or sops-nix
#       model = "gpt-4";
#     };
#
#     # Vector Storage Configuration
#     vectorStore = {
#       provider = "qdrant";  # or "chroma", "pinecone", etc.
#       config = {
#         host = "localhost";
#         port = 6333;
#       };
#     };
#
#     # Optional: Environment variables
#     environmentFile = "/etc/mem0/environment";
#   };
#
# Using with m3ta.ports (recommended):
#
#   m3ta.ports = {
#     enable = true;
#     definitions = { mem0 = 8000; };
#     hostOverrides.laptop = { mem0 = 8080; };
#     currentHost = config.networking.hostName;
#   };
#
#   m3ta.mem0 = {
#     enable = true;
#     port = config.m3ta.ports.get "mem0";  # Automatically uses host-specific port
#   };
{
  config,
  lib,
  pkgs,
  ...
}:
with lib; let
  cfg = config.m3ta.mem0;

  # Python environment with mem0
  pythonEnv = pkgs.python3.withPackages (ps:
    with ps; [
      cfg.package
    ]);

  # Convert vector store config to environment variables
  vectorStoreEnv =
    if cfg.vectorStore.provider == "qdrant"
    then {
      MEM0_VECTOR_PROVIDER = "qdrant";
      QDRANT_HOST = cfg.vectorStore.config.host or "localhost";
      QDRANT_PORT = toString (cfg.vectorStore.config.port or 6333);
      QDRANT_COLLECTION = cfg.vectorStore.config.collection_name or "mem0_memories";
    }
    else if cfg.vectorStore.provider == "pgvector"
    then {
      MEM0_VECTOR_PROVIDER = "pgvector";
      POSTGRES_HOST = cfg.vectorStore.config.host or "localhost";
      POSTGRES_PORT = toString (cfg.vectorStore.config.port or 5432);
      POSTGRES_DB = cfg.vectorStore.config.dbname or "postgres";
      POSTGRES_USER = cfg.vectorStore.config.user or "postgres";
      POSTGRES_PASSWORD = cfg.vectorStore.config.password or "postgres";
      POSTGRES_COLLECTION = cfg.vectorStore.config.collection_name or "mem0_memories";
    }
    else if cfg.vectorStore.provider == "chroma"
    then {
      MEM0_VECTOR_PROVIDER = "chroma";
      CHROMA_HOST = cfg.vectorStore.config.host or "localhost";
      CHROMA_PORT = toString (cfg.vectorStore.config.port or 8000);
      CHROMA_COLLECTION = cfg.vectorStore.config.collection_name or "mem0_memories";
    }
    else {};

  # Start script that sets up environment and runs the server
  startScript = pkgs.writeShellScript "mem0-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
    ''}

    # Load API key from file if specified
    ${optionalString (cfg.llm.apiKeyFile != null) ''
      if [ -f "${cfg.llm.apiKeyFile}" ]; then
        export OPENAI_API_KEY="$(cat ${cfg.llm.apiKeyFile})"
      fi
    ''}

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

    # Run the server
    exec ${pythonEnv}/bin/mem0-server
  '';
in {
  options.m3ta.mem0 = {
    enable = mkEnableOption "Mem0 REST API server";

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

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

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

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

    logLevel = mkOption {
      type = types.enum ["critical" "error" "warning" "info" "debug" "trace"];
      default = "info";
      description = "Logging level for the server.";
    };

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

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

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

    environmentFile = mkOption {
      type = types.nullOr types.path;
      default = null;
      description = ''
        Environment file containing additional configuration.
        This file should contain KEY=value pairs, one per line.
        Useful for secrets that shouldn't be in the Nix store.
      '';
      example = "/etc/mem0/environment";
    };

    # LLM Configuration
    llm = {
      provider = mkOption {
        type = types.enum ["openai" "anthropic" "azure" "groq" "together" "ollama" "litellm"];
        default = "openai";
        description = "LLM provider to use for memory operations.";
      };

      model = mkOption {
        type = types.str;
        default = "gpt-4o-mini";
        description = "Model name to use for the LLM.";
      };

      apiKeyFile = mkOption {
        type = types.nullOr types.path;
        default = null;
        description = ''
          Path to file containing the API key for the LLM provider.
          The file should contain only the API key.
          This is more secure than putting the key in the Nix store.
        '';
        example = "/run/secrets/openai-api-key";
      };

      temperature = mkOption {
        type = types.nullOr types.float;
        default = null;
        description = "Temperature parameter for LLM generation.";
      };

      maxTokens = mkOption {
        type = types.nullOr types.int;
        default = null;
        description = "Maximum tokens for LLM generation.";
      };

      extraConfig = mkOption {
        type = types.attrs;
        default = {};
        description = "Additional LLM configuration options.";
        example = {
          top_p = 1.0;
          frequency_penalty = 0.0;
        };
      };
    };

    # Vector Store Configuration
    vectorStore = {
      provider = mkOption {
        type = types.enum [
          "qdrant"
          "chroma"
          "pinecone"
          "weaviate"
          "faiss"
          "pgvector"
          "redis"
          "elasticsearch"
          "milvus"
        ];
        default = "qdrant";
        description = "Vector database provider to use.";
      };

      config = mkOption {
        type = types.attrs;
        default = {};
        description = ''
          Configuration for the vector store.
          The structure depends on the provider.
        '';
        example = literalExpression ''
          {
            host = "localhost";
            port = 6333;
            collection_name = "mem0_memories";
          }
        '';
      };
    };

    # Embedder Configuration
    embedder = {
      provider = mkOption {
        type = types.nullOr (types.enum ["openai" "huggingface" "ollama" "vertexai"]);
        default = null;
        description = "Embedding model provider. If null, uses default.";
      };

      model = mkOption {
        type = types.nullOr types.str;
        default = null;
        description = "Embedding model name to use.";
        example = "text-embedding-3-small";
      };

      config = mkOption {
        type = types.attrs;
        default = {};
        description = "Configuration for the embedder.";
        example = {
          model = "text-embedding-3-small";
        };
      };
    };
  };

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

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

    # Systemd service
    systemd.services.mem0 = {
      description = "Mem0 REST 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";
          MEM0_HOST = cfg.host;
          MEM0_PORT = toString cfg.port;
          MEM0_LLM_PROVIDER = cfg.llm.provider;
          MEM0_LLM_MODEL = cfg.llm.model;
          MEM0_HISTORY_DB_PATH = "${cfg.stateDir}/history.db";
          MEM0_WORKERS = toString cfg.workers;
          MEM0_LOG_LEVEL = cfg.logLevel;
        }
        // optionalAttrs (cfg.llm.temperature != null) {
          MEM0_LLM_TEMPERATURE = toString cfg.llm.temperature;
        }
        // optionalAttrs (cfg.llm.extraConfig != {}) {
          MEM0_LLM_EXTRA_CONFIG = builtins.toJSON cfg.llm.extraConfig;
        }
        // optionalAttrs (cfg.embedder.provider != null) {
          MEM0_EMBEDDER_PROVIDER = cfg.embedder.provider;
        }
        // optionalAttrs (cfg.embedder.model != null) {
          MEM0_EMBEDDER_MODEL = cfg.embedder.model;
        }
        // vectorStoreEnv;
    };

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